Fix concretization error with normalize_function

[teddykoker]

e3nn_jax.activation.normalize_function currently prevents shard map from working due to the .item() call. Reproducer:

import os
os.environ.setdefault("JAX_PLATFORMS", "cpu")
os.environ["XLA_FLAGS"] = " --xla_force_host_platform_device_count=2"

from functools import partial
import jax, jax.numpy as jnp
from jax import lax, pmap
from e3nn_jax._src.activation import normalize_function

mesh = jax.sharding.Mesh(jax.devices()[:2], ("i",))
P = jax.sharding.PartitionSpec


def foo(x):
    act = normalize_function(jax.nn.silu)
    return jnp.sum(act(x))

# pmap: works
@partial(pmap, axis_name="i")
def pmap_ok(x):
    return foo(lax.all_gather(x, "i").reshape(-1))

print("pmap:", pmap_ok(jnp.ones((2, 4))))

# shard_map: fails
@partial(jax.shard_map, mesh=mesh, in_specs=(P("i"),), out_specs=P(), check_vma=False)
def shard_map_fail(x):
    return foo(lax.all_gather(x, "i").reshape(-1))

try:
    print("shard_map:", shard_map_fail(jax.device_put(jnp.ones((2, 4)), jax.sharding.NamedSharding(mesh, P("i")))))
except jax.errors.ConcretizationTypeError as e:
    print(f"shard_map FAILS: {type(e).__name__}: {str(e)[:120]}...")

This change fixes.

[ameya98]

should we also be handling the case where c is close to 0?

[teddykoker]

this would only be the case if the activation function was lambda x: 0 or close... figured it wasn't an issue since it wasn't in the original code either

[teddykoker]

Not sure if it is an issue actually, e.g. atomicarchitects/nequix#38 works fine with shard_map.