All notable changes to this project will be documented in this file.
The format is based on Keep a Changelog, and this project adheres to Semantic Versioning.
- Fixed bug in
e3nn.SphericalSignal.to_s2point()with batched coefficients and points.
- Tests for GaussLegendre quadrature in
e3nn.SphericalSignal.
- Docstring for GaussLegendre quadrature in
e3nn SphericalSignal.
- Support for custom instructions and initializers in e3nn.flax.Linear and e3nn.haiku.Linear.
- Fix documentation build errors.
e3nn.wherefunction- Add optional
maskargument ine3nn.flax.BatchNorm
- replace
jnp.ndarraybyjax.Array
e3nn.onesande3nn.ones_likefunctionse3nn.equinoxsubmodule
- python 3.9 compatibility
- Fix
pyproject.toml, the documentation build was broken. Thanks to @SauravMaheshkar!
- Support for
s2fftine3nn.to_s2gridande3nn.from_s2grid, thanks to @ameya98! - Add a special case implementation for
e3nn.scatter_meanwhenmap_back and nel is not None.
e3nn.flax.BatchNorme3nn.scatter_mean- Add
e3nn.utils.vmapalso directly toe3nnmodule:e3nn.vmap
with_biasargument toe3nn.haiku.MultiLayerPerceptronande3nn.flax.MultiLayerPerceptron
- Improve compilation speed and stability of
s2gridfor largelmax(useis_normalized=Trueinlpmn_values)
- Add back the optimizations with the lazy
._chunksthat was removed in 0.19.0
e3nn.Irreps.mul_gcde3nn.IrrepsArray.extend_with_zerosto extend an array with zeros, can be useful for residual connections
- rewrite
e3nn.tensor_squareto be simpler (and faster?) - use
jax.scipy.special.lpmn_valuesto implemente3nn.legendre. Faster on GPU and supports reverse-mode differentiation. - [BREAKING] Change the output format of
e3nn.legendre!
- Add back a lazy
._chunksine3nn.IrrepsArrayto fix issue #38
- Fix missing support for zero flags in
e3nn.elementwise_tensor_product
- [BREAKING] Move
Instruction,FunctionalTensorProductandFunctionalFullyConnectedTensorProductintoe3nn.legacysubmodule - Reimplement
e3nn.tensor_productande3nn.elementwise_tensor_productin a simpler way
e3nn.utils.vmapto propagatezero_flagsin the vectorized function.
- Simplify the tetris examples
- Example of what is fixed: assume
x.ndim = 2, allowx[:, None]but preventx[:, :, None]andx[..., None]
- [BREAKING]
e3nn.flax.Linearande3nn.haiku.Linearnow don't output the impossible irreps anymore. To force the output of all irreps, useforce_irreps_out = True. For instancee3nn.flax.Linear("0e + 1o")("0e")will now return"0e"instead of"0e + 1o". - [BREAKING]
e3nn.utils.assert_equivarianthas the same signature ase3nn.utils.equivariance_test - [BREAKING] Move
as_irreps_array,zerosandzeros_likefrome3nn.IrrepsArraytoe3nn - [BREAKING] Move
IrrepsArray.from_listtoe3nn.from_chunks - [BREAKING] Rename
IrrepsArray.listintoIrrepsArray.chunks - [BREAKING] Rename
IrrepsArray.remove_nonesintoIrrepsArray.remove_zero_chunks e3nn.IrrepsArrayhas now only.arrayas data attribute.
e3nn.IrrepsArray.rechunke3nn.IrrepsArray.zero_flagsa tuple of bools that indicates which chunks are zero
- [BREAKING] Renamed
e3nn.utiline3nn.utils
Irreps.set_mul(int)to set the multiplicity of all irrepsIrreps.filter(lmax=int)to filter out irreps withl > lmaxIrrepsArray.filter(lmax=int)to filter out irreps withl > lmaxIrrepsArray.__radd__andIrrepsArray.__rsub__to supportscalar + IrrepsArrayandscalar - IrrepsArray0 + IrrepsArrayand0 - IrrepsArrayare now always accepted as special cases.- Support for
IrrepsArray / array - Add
utilsas a submodule
e3nn.scatteroperation handle indices withndim > 1
e3nn.crossfor completeness
- Optimize
e3nn.reduced_symmetric_tensor_product_basis, especially for thekeep_irargument
LinearSHTPmodule implementing the optimized linear mixing of inputs tensor product with spherical harmonicsD_from_axis_angleto_s2grid:quadrature="gausslegendre"by defaultsoft_oddactivation function for odd scalars- more support of arrays implicitely converted into
IrrepsArrayas scalars (i.e. added fewIrrepsArray.as_irreps_array)
scalar_activationsimpler to use with default activation functions (a bit like gate)
e3nn.normalize_functionnow uses a deterministic (not pseudorandom) algorithm to compute the normalization factor.
normalize_actoption toe3nn.scalar_activationande3nn.gate. We can now turn the normalization off if we want to.e3nn.norm_activationas a new activation function.
- Fix
NaNin the gradients ofe3nn.xyz_to_angles. The gradients are now0when the input is on the poles.
e3nn.dot: compute the dot product between twoIrrepsArrayper_irrepargument toe3nn.norm: compute the norm of each irrep independently ifper_irrep=Truee3nn.tensor_product_with_spherical_harmonicsfrom https://arxiv.org/pdf/2302.03655.pdf
__repr__(Irreps())has been changed from""to"Irreps()"
- spherical harmonics edge case when
output_irreps=Irreps()
e3nn.SphericalSignal.sampleto sample a point on the spheree3nn.scatter_max
- [BREAKING] Removed
e3nn.s2_sum_of_diracsin favor ofe3nn.s2_dirac - [BREAKING]
e3nn.gradnow regroups the output by default. It can be disabled withregroup_output=False
e3nn.SphericalSignalarithmetic operationse3nn.Irreps.D_from_anglescomputes (again!) the Wigner D matrices using the J matrices for L <= 11. This is faster and more accurate than using the expm.
e3nn.SphericalSignalclass to represent signals on the sphereSignal on the Spheresection in the documentatione3nn.Irreps.D_from_log_coordinatesrotation_angle_from_*functionse3nn.to_s2pointfunction
- Wigner D matrices are computed from the log coordinates which makes 1 instead of 3 calls to
expm. - [BREAKING]
e3nn.util.assert_output_dtyperenamed toe3nn.util.assert_output_dtype_matches_input_dtype - [BREAKING] Update
experimental.point_convolutionto use the last changes. - [BREAKING] changed the
e3nn.to_s2gridande3nn.from_s2gridsignature and default normalization.
- [BREAKING] All the
haikumodules from the main module. They are now in thee3nn.haikusubmodule. - [BREAKING]
e3nn.wigner_Din favor ofe3nn.Irrep.D_from_*
- Removed
jax.jitdecorator toIrreps.D_from_*that was causing a bug.
e3nn.s2grid_vectorsande3nn.pad_to_plot_on_s2gridto help plotting signals on the spheree3nn.util.assert_output_dtypeto check the output dtype of a functione3nn.s2_irrepsis a function to create the irreps of the coefficients of a signal on the spheree3nn.reduced_antisymmetric_tensor_product_basisto compute the basis of the reduced antisymmetric tensor productIrrepsArray * scalaris supported if the number of scalars matches the number of irreps
- Optimize the
reduced_symmetric_tensor_product. It is now up to 100x faster than the previous implementation. e3nn.from_s2gridande3nn.to_s2gridare now more flexible with input and output irreps, you can skip some l's and have them in any order- [BREAKING]
e3nn.from_s2gridrequires andirrepsargument instead of almaxargument
- Increase robusteness of
e3nn.spherical_harmonicstowardsnanwhennormalize=True
IrrepsArray.astypeto cast the underlying arraye3nn.flax.MultiLayerPerceptronande3nn.haiku.MultiLayerPerceptrone3nn.IrrepsArray.from_list(..., dtype)- Add sparse tensor product as an option in
e3nn.tensor_productand related functions. It sparsify the clebsch gordan coefficients. It has more inpact whenfused=True. It is disabled by default because no improvement was observed in the benchmarks. - Add
log_coordinatesalong the other parameterizations of SO(3).e3nn.log_coordinates_to_matrix,e3nn.rand_log_coordinates, etc.
- set dtype for all
jnp.zeros(..., dtype)calls in the codebase - set dtype for all
jnp.ones(..., dtype)calls in the codebase
- [BREAKING]
e3nn.full_tensor_productin favor ofe3nn.tensor_product - [BREAKING]
e3nn.FunctionalTensorSquarein favor ofe3nn.tensor_square - [BREAKING]
e3nn.TensorSquarein favor ofe3nn.tensor_square - [BREAKING]
e3nn.IrrepsArray.catin favor ofe3nn.concatenate - [BREAKING]
e3nn.IrrepsArray.randnin favor ofe3nn.normal - [BREAKING]
e3nn.Irreps.randnin favor ofe3nn.normal - [BREAKING]
e3nn.Irreps.transform_by_*in favor ofe3nn.IrrepsArray.transform_by_*
- moves
BatchNormandDropouttoe3nn.haikusubmodule, will remove them from the main module in the future. - move
e3nn.haiku.FullyConnectedTensorProductinhaikusubmodule. Undeprecate it because it's faster thane3nn.tensor_productfollowed bye3nn.Linear. This is becauseopteinsumoptimizes the contraction of the two operations.
e3nn.scatter_sumto replacee3nn.index_add.e3nn.index_addis deprecated.- add
flaxandhaikusubmodules. Plan to migrate all modules toflaxandhaikuin the future. - Implement
e3nn.flax.Linearand movee3nn.Linearine3nn.haiku.Linear.
- [BREAKING]
3 * e3nn.Irreps("0e + 1o")now returns3x0e + 3x1oinstead of1x0e + 1x1o + 1x0e + 1x1o + 1x0e + 1x1o - [BREAKING] in Linear, renamed
num_weightstonum_indexed_weightsbecause it was confusing.
e3nn.Irreps("3x0e + 6x1o") // 3returns1x0e + 2x1o
s2gridis now jitable
e3nn.Irreps.regroupande3nn.IrrepsArray.regroupto regroup irreps. Equivalent tosortfollowed bysimplify.- add
regroup_outputparameter toe3nn.tensor_productande3nn.tensor_squareto regroup the output irreps.
e3nn.IrrepsArray.convertis now private (e3nn.IrrepsArray._convert) because it's recommended to other methods instead.- breaking change use
input.regroup()ine3nn.Linearwhich can change the structure of the parameters dictionary. - breaking change
regroup_outputisTrueby default ine3nn.tensor_productande3nn.tensor_square. - To facilitate debugging, if not
keyis provided toe3nn.normalit will use the hash of the irreps. - breaking change changed normalization of
e3nn.tensor_squarein the case ofnormalized_input=True
- Deprecate
e3nn.TensorSquare
e3nn.Linearnow supports integer "weights" inputs.e3nn.Linearnow supportsnameargument.- Add
.dtypetoIrrepsArrayto get the dtype of the underlying array.
e3nn.MultiLayerPerceptronnames its layerslinear_0,linear_1, etc.
- s2grid:
e3nn.from_s2gridande3nn.to_s2gridthanks to @songk42 for the contribution - argument
max_order: intto functionreduced_tensor_product_basisto be able to limit the polynomial order of the basis MultiLayerPerceptronacceptsIrrepsArrayas input and outpute3nn.Linearaccepts optional weights as arguments that will be internally mixed with the free parameters. Very usefyul to implement the depthwise convolution
- breaking change
e3nn.normalhas a new argument to get normalized vectors. - breaking change
e3nn.tensor_squarenow distinguishes betweennormalization=normandnormalized_input=True.
e3nn.SymmetricTensorProductoperation: a parameterized version ofx + x^2 + x^3 + ....e3nn.soft_envelopea smoothC^infenvelope radial function.e3nn.tensor_square
Irrep.generatorsandIrreps.generatorsfunctions to get the generators of the representations.e3nn.besselfunctionslice_by_mul,slice_by_dimandslice_by_chunkfunctions toIrrepsandIrrepsArray
- breaking change
e3nn.soft_one_hot_linspacedoes not supportbesselanymore. Usee3nn.besselinstead. e3nn.gateis now more flexible of the input format, see examples in the docstring.
- breaking change
IrrepsArray.split
- fix
IrrepsArray.zeros().at[...].add
e3nn.reduced_symmetric_tensor_product_basis(irreps: Irreps, order: int)e3nn.IrrepsArray.filtered(keep: List[Irrep])e3nn.reduced_tensor_product_basis(formula_or_irreps_list: Union[str, List[e3nn.Irreps]], ...)IrrepsArray.at[i].set(v)andIrrepsArray.at[i].add(v)- add
Irreps.is_scalar
- Simple irreps indexing of
IrrepsArray: likex[..., "10x0e"]but notx[..., "0e + 1e"] e3nn.concatenate, e3nn.mean, e3nn.sume3nn.normforIrrepsArraye3nn.tensor_producte3nn.normal- Better support of
+ - * /operators forIrrepsArray - Add new operator
e3nn.grad: it takes anIrrepsArray -> IrrepsArrayfunction and returns aIrrepsArray -> IrrepsArrayfunction - Add support of operator
IrrepsArray ** scalar - Add support of
x[..., 3:6]forIrrepsArray - Add
e3nn.reduced_tensor_product_basis - Add
e3nn.stack
IrrepsArray.catis now deprecated and replaced bye3nn.concatenatee3nn.full_tensor_productis now deprecated and replaced bye3nn.tensor_producte3nn.FullyConnectedTensorProductis now deprecated in favor ofe3nn.tensor_productande3nn.Linear- breaking change remove
IrrepsArray.from_any - breaking change remove option
optimize_einsums, (it is now alwaysTrue)
- breaking change rewrite the
equivariance_errorandassert_equivariantfunctions
- breaking change change the ordering of
Irrep. Now it matches withIrrep.iterator. - breaking change
Irrep("1e") == "1e"andIrreps("1e + 2e") == "1e + 2e"are nowTrue. - breaking change
Linearsimplify theirreps_outwhich might cause reshape of the parameters. index_addsupportsIrrepArray
- broadcast for
Linear - argument
channel_outtoLinearfor convenience Irrepscan be created from aMulIrrep"0e" + Irreps("1e")is now supported"0e" + Irrep("1e")is now supportedmap_backargument toindex_addIrrepsArray.split(list of irreps)poly_envelopefunction
- breaking change rename
IrrepsDataintoIrrepsArray - breaking change
IrrepsArray.shapeis now equal tocontiguous.shape(instead ofcontiguous.shape[:-1]) - breaking change
IrrepsArray * arrayrequiresarray.shape[-1]to be 1 orarrayto be a scalar - breaking change
IrrepsArray.contiguousis renamed inIrrepsArray.array - breaking change
IrrepsArray.newis renamed inIrrepsArray.from_any spherical_harmonicsnormalization is now set tocomponentlike everything else.
- breaking change
IrrepsArray.from_contiguousis removed. UseIrrepsArray(irreps, array)instead.
- add
e3nn.configto set global defaults parameters __getindex__toIrrepsDatagradient_normalizationargument that can beelementorpathpath_normalizationcan be a number between 0 and 1- add nearest interpolation for
zoom, default is linear - implement
custom_jvpfor spherical harmonics
- Docker image
- add the
shfunction that does not useIrrepsDataas input/output legendrealgorithm to compute spherical harmonics- add flag
algorithmto specify the algorithm to use for computing spherical harmonics, uselegendrefor large L. experimental.voxel_convolution: add optional dynamic steps (not static for jit)
- fix a bug in
experimental.voxel_convolutionconstructor
- Function
matrixtoFunctionalLinear experimental.voxel_convolution:paddingand add self-connection into the convolution kernelexperimental.voxel_pooling: addoutput_sizeargument to thezoomfunctionIrrepsData:listattribute is now lazily initializedexperimental.voxel_convolution: add possibility to have different radial functions depenfing on the spherical harmonic degree
- Behavior of
epsinBatchNorm. Nowinput / sqrt((1 - eps) * norm^2 + eps)instead ofinput / sqrt(norm^2 + eps) - Optimized
spherical_harmonicsby decomposing the order in powers of 2. It is supposed to improve stability because less operations are performed for high orders. It improves the performance when computing a single order. - Optimized
spherical_harmonicsby using dense matrix multiplication instead of sparse matrix multiplication.
- add
loopargument toradius_graph
- use
dataclasses.dataclassinstead of customdataclass - Get Clebsch-Gordan coefficients from qutip and a change of basis
- Add
start_zeroandend_zeroarguments to functionsoft_one_hot_linspace
IrrepsDatacan be given as argument ofspherical_harmonics- added broadcasting of
IrrepsData,elementwise_tensor_product,FullyConnectedTensorProduct,full_tensor_product
BatchNormsupports NoneBatchNormsupports change default value ofepsfrom1e-5to1e-4gatechange default odd activation to (1 - exp(x^2)) * x
gatelist of activations argument is now optionalexperimental.transformer.Transformersimplified interface usingIrrepsDataand swap two arguments order
IrrepsData.repeat_irreps_by_last_axisIrrepsData.repeat_mul_by_last_axisIrrepsData.factor_mul_to_last_axis- add
axisargument toIrrepsData.cat IrrepsData.remove_nonesIrrepsData.ones
experimental.point_convolution.Convolutionsimplified interface usingIrrepsData
- Changelog