# pytensor hackathon planning

Goals:

- Familiarize as many people as possible with the codebase
- Figure out important topics for future work

## Update docs

A lot of documentation is using aesara directly, this needs to be changed
systematically.

## Rename aesara -> pytensor

+ [x] Renames of project names in the repository
+ [ ] `import pytensor.tensor as at` → `import pytensor.tensor as pt` (low-priority)

## Upgrading PyMC
+ [ ] Bump to Aesara 2.8.9 in PyMC
+ [ ] Locally test PyMC main with PyTensor main
+ [ ] If needed, make compatibility changes in PyMC
+ [ ] Swap PyMC to depend on PyTensor

## Governance & CoC updates
+ [x] Have been updated to refer to PyMC governance & CoC

## Maintenance setup
+ [ ] conda-forge feedstock ⏳ https://github.com/conda-forge/staged-recipes/pull/21253
+ [x] readthedocs
+ [x] release automation

## Improve numba backend coverage and testing.

Known issues with numba backend:

- Compilation is often not tested.
- Missing special functions from scipy. Some prior work [here](https://github.com/aesara-devs/aesara/pull/1078).
- Missing scipy/numpy linalg functions.
- scan performance is sometimes lacking
- numba graphs involving random draws? (ie RandomState)
- Ricardo: RandomState support is pretty lame. RandomGenerator support seems to be getting better. PR here: https://github.com/aesara-devs/aesara/pull/1245
- Compile time can be long

## Can we remove the c backend?

Not currently, the numba backend isn't a full replacement yet,
but do we want to do this in the long term?

This would simplify distributing aesara a great deal, because
we wouldn't have to worry about a compiler in the user environment
anymore.

Note: C is the most flexible backend, but the hardest to maintain.

Max: i think we can at least rewrite c backend in the same dispatch manner as we have for numba or jax. it will bring consistency to the codebase. Same could be done for python backend

## Jax backend

Scan is broken: Remi has been working on it in Aesara: https://github.com/aesara-devs/aesara/pull/1202, https://github.com/aesara-devs/aesara/pull/1284

Random graphs are represented completely different in Aesara and JAX, there are some hacks to make it kind of work: https://github.com/aesara-devs/aesara/pull/1214

No windows support

## Merge aeppl back into pymc

We decided to move aeppl back into pymc.

## Dynamic or static broadcasting

Recently aesara switch to allow dynamic broadcasting.
Some background can be found [here](https://github.com/aesara-devs/aesara/discussions/1273).

Do we want to support that, and if so, how?

We need to deal with issues that were created by that change,
either by fixing some current bugs or reverting the original change
to allow dynamic broadcasting.

(This is a longer topic, needs a separate writeup I think)

-> Move back to static shapes, use shape inference for broadcastable info

## Improve static shape inference

Currently, there are two different systems in place that
figure out shapes of tensors: Static shape inference in
the `Op.make_node` methods, and separate infer_shape functions.

We probably want to merge those.
(This might interact with support for dimensions, as that probably
also needs changes to `make_node`)

Ricardo: Note: infer_shape always works, it falls-back to executing the Operation and taking the shape if needed. It also returns symbolic expressions, whereas for `make_node` we need concrete values. Constant-folding may be too slow in practice (although if coverage improves, it may not be so bad, because we always have static shapes for the inputs)

## Rewrite safety

Theano had lots of unsafe rewrites (ie rewrites that assume that
the input of a theano function is valid). There were a lot of
changes to change those, so that errors are not optimized away.
There are at least a few of those left though, and the dynamic
shape changes might have also had an impact on this.

It might actually be interesting to change the framework of
assumptions of individual ops. Could Asserts be properties
of ops or variables, so that they interfere less with rewrites?

## Convert some ops to OpFromGraph

Some ops (eg softmax) only need to be custom ops so that we
control their gradient. Do we want to convert them to OpFromGraph,
so that backends need to implement fewer ops?

Ricardo: Yes, but we may want to create an ElemwiseOpFromGraph. Right now OpFromGraph must be compiled with specific input types, which is wasteful.

## Support for dimensions

aesara currently only knows about integer shapes, but it would
be great if it could make use of dimension information.

Do we want to support that? How? Including xarray style broadcasting?

Do we want to have ordered dims?
Should each dimension have a name?
Should that name be unique?

country = pt.Dim("country")
age = pt.Dim("age")
countr_effect = pt.dims.dvector(dims=country)
age_effect = pt.dims.dvector(dims=age)
effect = country_effect + age_effect   # dims=[country, age]

# Go from new api to old
effect_old_api = effect.to_shaped([country, age])

effect.ravel()  # dims={ProductDim({country, age})}

effect.rename_dim(country) + effect  # dims=(RenamedDim(country), country, age)

# Go from old api to new
second_dim = pt.Dim()  # Do we have to specify a shape?
effects_imported = pt.dims.from_shaped(effect_old_api, dims=[county, second_dim])

effects_imported.sum(second_dim)

effects.sel(country=["USA", "Italy"])  # dims=(SelDim(country, ["USA", "Italy"]), age)

## (Semi-)Automatic marginalization

Ricardo has been experimenting with this.

## Merge Composite Elemwise ops where possible

It should often be possible to avoid looping over data
several times, if we merge different elemwise ops that use
the same inputs. I think this has the potential to speed up
a lot logp graphs considerably, especially if combined with
the next item. POC implementation (tests still failing) here: https://github.com/aesara-devs/aesara/pull/1242

dist = pm.Normal.dist(mu=np.ones(10), sigma=np.ones(10))
logp = pm.logp(dist, x)

grad = pt.grad(logp, x)

func = pt.function([x], [logp, grad])


## we want
logp = 0.
grad = np.empty_like(x)
for i in range(len(x)):
    dist = -x[i] ** 2
    logp -= dist
    grad[i] = -x

# right now

y = np.empty_like(x)
for i in range(len(x)):
    y[i] = -val ** 2
    
logp = sum(y)

grad = np.empty_like(x)
for i in range(len(x)):
    grad[i] = -x

Need op
ElemwiseSum(scalar_op=Composite([in1, in2], [out1, out2]), sum_axes=[False, True])

## Optimize sum(elemwise) ops

We often compute sums of elemwise ops, which is slower than
just summing while computing the elemwise. We could do this
by introducing a new SumElemwise op, that does both at the same time. Brandon started something here: https://github.com/aesara-devs/aesara/pull/1285

## review clone_replace / scan
these operations are essential to work on pymc and manipulating graph.

clone_replace: reading the code i have an inression it is overly complicated and has to be refactored to gain readability and improve overall maintainance.

scan: it has some undocumented limitations, strict mode, which if you do not use code can break later without any meaningfull error. scan is essential op. we can improve that subclassing ops into more specific Map/vectorize ops which can be later optimized in other backends i think improving code quality there will bring huge gains as well

easy scan things:
- improve docs
- improve error messages

others:
- Add specialized scan ops (reduce, elemwise (see blockwise PR in aesara))
- Break up scan into other ops

Start adding simpler scan-like ops that only do something simpler. (bottom up)

## Think about a mlir backend

See https://mlir.llvm.org/ which also comes with a python interface.

## Think about vmap similar things

Blockwise (aesara PR)
How does this interact with dims?