[nnx] fix Variable overloads and add shape/dtype properties

[cgarciae]

What does this PR do?

• Fixes Variable inplace operators.
• Updates the nnx_basics guide to use inplace updates on the Variable
• Adds .shape and .dtype properties to Variable.

After these changes you can now correctly do in-place operations of Variables:

class Count(nnx.Variable): ...

class Counter(nnx.Module):
  def __init__(self):
    self.count = Count(jnp.array(0, dtype=jnp.uint32))
    
  def increment(self):
    self.count += 1

Previously you have do this on the value: self.count.value += 1 .

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB