mul_ir_dict

cuequivariance_jax.ir_dict.mul_ir_dict(irreps, data)

Create a dict[Irrep, data] by broadcasting data to match irreps structure.

Useful for creating output templates or broadcasting scalar values across irreps.

Parameters:

• irreps (Irreps) – Irreps specification defining the dict keys.

• data (Any) – Data to broadcast to each irrep key.

Returns:

Dictionary with irrep keys, each value set to data.

Return type:

dict[Irrep, Any]

Example

>>> import cuequivariance as cue
>>> irreps = cue.Irreps(cue.O3, "128x0e + 64x1o")
>>> batch = 32
>>> template = mul_ir_dict(irreps, jax.ShapeDtypeStruct((batch,), jnp.float32))