Repository: google/flax

Issue:
np.array parameters may lead to a silent failure
Passing np.array parameters (instead of jnp.array) to a linen module may lead to a silent failure, see the following example:

```
import flax.linen as nn

import jax
import jax.numpy as jnp

import numpy as np

t = jnp.zeros([2, 196], jnp.int32)
print(f'Input shape: {t.shape}')

m = nn.Embed(32, 10)
rng = jax.random.PRNGKey(0)
vars = m.init(rng, t)
o1 = m.apply(vars, t)
print(f'Expected output shape: {o1.shape}')

o2 = m.apply(jax.tree_map(np.array, vars), t)
print(f'Numpy output shape: {o2.shape}')
```

Output:

```
Input shape: (2, 196)
Expected output shape: (2, 196, 10)
"Numpy params" output shape: (196,)  <-- Different output shape
```



Discussion:
Thanks for catching this!

When you map the embedding to an `np.array`, what will happen when applying the `Embed` module is that the embedding (which is now a Numpy array) is indexed with a `jax.numpy` array. This causes Numpy to treat the `jnp.array` as a tuple, which is not what we want:

```python
embedding = jnp.array([[1], [2]])
idx = jnp.array([0], jnp.int32)

embedding[idx]  # Similar to your o1 -- correct
>>> [[1]]

np.array(embedding)[idx]  # Similar to o2 -- wrong
>>> [1]
```

We can verify the `jnp` is cast incorrectly by changing its type:

```python
np.array(embedding)[tuple(idx)]
>>> [1]

np.array(embedding)[np.array(idx)]
>>> [[1]]
```

Actually this also throws a Deprecationwarning, suggesting we should explicitly cast the index to either np.arr or tuple.

So I think we can do two things:

1. Force the embedding to be a jnp array (and not a numpy array)
2. Check whether the embedding is a np array, and if so, explicitly cast the indexer to np.array.
I think we want to use `jnp.asarray` here to avoid unnecessary copies but guarantee that the embedding is a jax numpy array.
@jheek I think that may still break if the params aren't jax numpy arrays, which can happen if you load weights from a file. See also #1261.

Resolving patch:
diff --git a/flax/linen/linear.py b/flax/linen/linear.py
--- a/flax/linen/linear.py
+++ b/flax/linen/linear.py
@@ -407,7 +407,8 @@ def __call__(self, inputs):
     """
     if not jnp.issubdtype(inputs.dtype, jnp.integer):
       raise ValueError('Input type must be an integer or unsigned integer.')
-    return self.embedding[inputs]
+    # Use take because fancy indexing numpy arrays with JAX indices does not work correctly.
+    return jnp.take(self.embedding, inputs, axis=0)
 
   def attend(self, query):
     """Attend over the embedding using a query array.


Test patch:
diff --git a/tests/linen/linen_linear_test.py b/tests/linen/linen_linear_test.py
--- a/tests/linen/linen_linear_test.py
+++ b/tests/linen/linen_linear_test.py
@@ -272,6 +272,21 @@ def test_embed(self):
     z = embed_module.apply(initial_params, jnp.ones((3,)), method=embed_module.attend)
     np.testing.assert_allclose(z, 3. * jnp.arange(4))
   
+  def test_embed_numpy(self):
+    rng = dict(params=random.PRNGKey(0))
+    x = jnp.arange(4)[None]
+    dummy_embedding = np.broadcast_to(
+        np.arange(4)[..., None], (4, 3)).astype(np.float32)
+    embed_module = nn.Embed(
+        num_embeddings=4,
+        features=3,
+        embedding_init=lambda rng, shape, dtype: dummy_embedding,
+    )
+    y, initial_params = embed_module.init_with_output(rng, x)
+    np.testing.assert_allclose(y, dummy_embedding[None])
+    z = embed_module.apply(initial_params, jnp.ones((3,)), method=embed_module.attend)
+    np.testing.assert_allclose(z, 3. * jnp.arange(4))
+  
   def test_non_final_axis(self):
     class Foo(nn.Module):
       @nn.compact