E3nn-jax

Added
- `e3nn.Irreps.regroup` and `e3nn.IrrepsArray.regroup` to regroup irreps. Equivalent to `sort` followed by `simplify`.
- add `regroup_output` parameter to `e3nn.tensor_product` and `e3nn.tensor_square` to regroup the output irreps.

Changed
- `e3nn.IrrepsArray.convert` is now private (`e3nn.IrrepsArray._convert`) because it's recommended to other methods instead.
- **breaking change** use `input.regroup()` in `e3nn.Linear` which can change the structure of the parameters dictionary.
- **breaking change** `regroup_output` is `True` by default in `e3nn.tensor_product` and `e3nn.tensor_square`.
- To facilitate debugging, if not `key` is provided to `e3nn.normal` it will use the hash of the irreps.
- **breaking change** changed normalization of `e3nn.tensor_square` in the case of `normalized_input=True`

Removed
- Deprecate `e3nn.TensorSquare`

Added
- `e3nn.Linear` now supports integer "weights" inputs.
- `e3nn.Linear` now supports `name` argument.
- Add `.dtype` to `IrrepsArray` to get the dtype of the underlying array.

Changed
- `e3nn.MultiLayerPerceptron` names its layers `linear_0`, `linear_1`, etc.

Added
- s2grid: `e3nn.from_s2grid` and `e3nn.to_s2grid` thanks to songk42 for the contribution
- argument `max_order: int` to function `reduced_tensor_product_basis` to be able to limit the polynomial order of the basis
- `MultiLayerPerceptron` accepts `IrrepsArray` as input and output
- `e3nn.Linear` accepts optional weights as arguments that will be internally mixed with the free parameters. Very usefyul to implement the depthwise convolution

Changed
- **breaking change** `e3nn.normal` has a new argument to get normalized vectors.
- **breaking change** `e3nn.tensor_square` now distinguishes between `normalization=norm` and `normalized_input=True`.

Added
- `e3nn.SymmetricTensorProduct` operation: a parameterized version of `x + x^2 + x^3 + ...`.
- `e3nn.soft_envelope` a smooth `C^inf` envelope radial function.
- `e3nn.tensor_square`

Added
- `Irrep.generators` and `Irreps.generators` functions to get the generators of the representations.
- `e3nn.bessel` function
- `slice_by_mul`, `slice_by_dim` and `slice_by_chunk` functions to `Irreps` and `IrrepsArray`

Changed
- **breaking change** `e3nn.soft_one_hot_linspace` does not support `bessel` anymore. Use `e3nn.bessel` instead.
- `e3nn.gate` is now more flexible of the input format, see examples in the docstring.

Removed
- **breaking change** `IrrepsArray.split`

Fixed
- fix `IrrepsArray.zeros().at[...].add`