import mx

# Create parameters
a = mx.Parameter(mx.randn([2, 3]))
b = mx.Parameter(mx.randn([2, 3]))

# Forward pass
c = a * b
loss = c.mean()

# Backward pass
loss.backward()

# Access gradients
print(a.grad)
print(b.grad)