SymPy Indexed Expressions and Implementation with Autodifferentiation

[Tareq Jaouni]

Hi!

I am currently working on a project that uses SymPy to calculate large-scale expansions of symbolic objects with non-integer indices (in general). We then convert the final expression into a numeric one, with the eventual goal of calculating a metric for optimization tasks. 

I would like to understand what happens under the hood when we add two indexed symbolic terms together. An example would be something like expanding the expression (a[0.0292*I (this is a complex number), 0, 0, 0] +  a[0.0292*I, 0, 0, 0])^2. I have read that expressions, such that those that are indexed, are stored in a tree-based representation. What does SymPy check for when adding, multiplying, etc different symbols?  Would there be a resource which would already discuss this? 

An ulterior motivation (and potentially a separate question) would be to combine our symbolic expansions with auto-differentiable frameworks like JAX. I am aware that there are packages like Sympy2Jax, but they don't seem to interface with Indexed expressions in my experience. What challenges might there be with adapting Indexed expressions to support this, particularly with very large expansions of indexed symbolic terms? 

Thanks! 

Tareq