E3nn-jax

Added
- `e3nn.SphericalSignal.sample` to sample a point on the sphere
- `e3nn.scatter_max`

Changed
- **[BREAKING]** Removed `e3nn.s2_sum_of_diracs` in favor of `e3nn.s2_dirac`
- **[BREAKING]** `e3nn.grad` now regroups the output by default. It can be disabled with `regroup_output=False`

Fixed
- `e3nn.SphericalSignal` arithmetic operations
- `e3nn.Irreps.D_from_angles` computes (again!) the Wigner D matrices using the J matrices for L <= 11. This is faster and more accurate than using the expm.

Added
- `e3nn.SphericalSignal` class to represent signals on the sphere
- `Signal on the Sphere` section in the documentation
- `e3nn.Irreps.D_from_log_coordinates`
- `rotation_angle_from_*` functions
- `e3nn.to_s2point` function

Changed
- Wigner D matrices are computed from the log coordinates which makes 1 instead of 3 calls to `expm`.
- **[BREAKING]** `e3nn.util.assert_output_dtype` renamed to `e3nn.util.assert_output_dtype_matches_input_dtype`
- **[BREAKING]** Update `experimental.point_convolution` to use the last changes.
- **[BREAKING]** changed the `e3nn.to_s2grid` and `e3nn.from_s2grid` signature and default normalization.

Removed
- **[BREAKING]** All the `haiku` modules from the main module. They are now in the `e3nn.haiku` submodule.
- **[BREAKING]** `e3nn.wigner_D` in favor of `e3nn.Irrep.D_from_*`

Fixed
- Removed `jax.jit` decorator to `Irreps.D_from_*` that was causing a bug.

Added
- `e3nn.s2grid_vectors` and `e3nn.pad_to_plot_on_s2grid` to help plotting signals on the sphere
- `e3nn.util.assert_output_dtype` to check the output dtype of a function
- `e3nn.s2_irreps` is a function to create the irreps of the coefficients of a signal on the sphere
- `e3nn.reduced_antisymmetric_tensor_product_basis` to compute the basis of the reduced antisymmetric tensor product
- `IrrepsArray * scalar` is supported if the number of scalars matches the number of irreps

Changed
- Optimize the `reduced_symmetric_tensor_product`. It is now up to 100x faster than the previous implementation.
- `e3nn.from_s2grid` and `e3nn.to_s2grid` are now more flexible with input and output irreps, you can skip some l's and have them in any order
- **[BREAKING]** `e3nn.from_s2grid` requires and `irreps` argument instead of a `lmax` argument

Fixed
- Increase robusteness of `e3nn.spherical_harmonics` towards `nan` when `normalize=True`

Added
- `IrrepsArray.astype` to cast the underlying array
- `e3nn.flax.MultiLayerPerceptron` and `e3nn.haiku.MultiLayerPerceptron`
- `e3nn.IrrepsArray.from_list(..., dtype)`
- Add sparse tensor product as an option in `e3nn.tensor_product` and related functions. It sparsify the clebsch gordan coefficients. It has more inpact when `fused=True`. It is disabled by default because no improvement was observed in the benchmarks.
- Add `log_coordinates` along the other parameterizations of SO(3). `e3nn.log_coordinates_to_matrix`, `e3nn.rand_log_coordinates`, etc.

Fixed
- set dtype for all `jnp.zeros(..., dtype)` calls in the codebase
- set dtype for all `jnp.ones(..., dtype)` calls in the codebase

Removed
- **[BREAKING]** `e3nn.full_tensor_product` in favor of `e3nn.tensor_product`
- **[BREAKING]** `e3nn.FunctionalTensorSquare` in favor of `e3nn.tensor_square`
- **[BREAKING]** `e3nn.TensorSquare` in favor of `e3nn.tensor_square`
- **[BREAKING]** `e3nn.IrrepsArray.cat` in favor of `e3nn.concatenate`
- **[BREAKING]** `e3nn.IrrepsArray.randn` in favor of `e3nn.normal`
- **[BREAKING]** `e3nn.Irreps.randn` in favor of `e3nn.normal`
- **[BREAKING]** `e3nn.Irreps.transform_by_*` in favor of `e3nn.IrrepsArray.transform_by_*`

Changed
- moves `BatchNorm` and `Dropout` to `e3nn.haiku` submodule, will remove them from the main module in the future.
- move `e3nn.haiku.FullyConnectedTensorProduct` in `haiku` submodule. Undeprecate it because it's faster than `e3nn.tensor_product` followed by `e3nn.Linear`. This is because `opteinsum` optimizes the contraction of the two operations.

Added
- `e3nn.scatter_sum` to replace `e3nn.index_add`. `e3nn.index_add` is deprecated.
- add `flax` and `haiku` submodules. Plan to migrate all modules to `flax` and `haiku` in the future.
- Implement `e3nn.flax.Linear` and move `e3nn.Linear` in `e3nn.haiku.Linear`.

Changed
- **[BREAKING]** `3 * e3nn.Irreps("0e + 1o")` now returns `3x0e + 3x1o` instead of `1x0e + 1x1o + 1x0e + 1x1o + 1x0e + 1x1o`
- **[BREAKING]** in Linear, renamed `num_weights` to `num_indexed_weights` because it was confusing.

Added
- `e3nn.Irreps("3x0e + 6x1o") // 3` returns `1x0e + 2x1o`

Fixed
- `s2grid` is now jitable