Support for Jax

[swamidass]

Excellent tool. We are likely to cite you in the future.

Would you consider building out a Jax backend?

[andreped]

Hello, @swamidass! I hope you are having a nice christmas holiday.

I've been considering getting into Jax, just never had the time. Mostly use TF myself, but use torch from time to time.

I might make an attempt to add Jax backend, if it is cruicial for a study of yours. Do you have a deadline when you need this feature to be ready?

I have made PRs to add full backend support for reinhard and a modified reinhard variant a few days ago. Could add Jax support, but would need to make an attempt first.

[swamidass]

Thanks for the response, and I hope you also are having a nice holiday.

We are in the middle of porting our WSI pipeline from tensorflow to Jax. So the there isn't a rush, because we are still using tenorflow. However, we would be well situated to test a jax backend.

[andreped]

We are in the middle of porting our WSI pipeline from tensorflow to Jax.

That's interesting. What is the main reason you are moving to Jax from TF? I was considering it, but for training purposes tf/keras offer most of what I need. For deployment, on the other hand, I find that using the inference engines such as ONNX Runtime or TensorRT are better, but that is not a problem as "all" TF models can be converted to the ONNX format, which both IEs support.

However, we would be well situated to test a jax backend.

I'm partly back from vacation now and will likely make an attempt at the Jax backend soon. Will keep you updated :]

[andreped]

It was too tempting...

Jax backend with Macenko seems to work. Did not spend lots of time with it. Hence, there are probably lots of ways to optimize it and whatnot, but at least we have a baseline that you could test and benchmark, if you'd like, @swamidass. Then we can try to improve it later on if necessary.

To test, be sure to install from this branch by:

pip install git+https://github.com/andreped/torchstain.git@jax-backend

Then to use, simply do:

import torchstain
import numpy as np

normalizer_jax = torchstain.normalizers.MacenkoNormalizer(backend='jax')
normalizer_jax.fit(target)
norm_jax, _, _ = normalizer_jax.normalize(to_transform)
norm = np.asarray(norm_jax)

When @carloalbertobarbano is back from vacation, we will resolve some existing PRs and then I can start adding Jax to the main branch for all stain normalization techniques. Should be quite seemless, as Jax has a surprisingly neat numpy-like API, when using jax.numpy as a direct replacement for numpy.

When all that is working, I believe we are ready for a new release :] Will keep you updated, @swamidass.

[swamidass]

Very nice! In our experience, JAX is substantially faster than TF and pytorch, which is why we are migrating to it.

So it will be interesting to see if that's true in this case too.

[swamidass]

One thing you want to be sure to test though, is whether or not the entrypoint you are providing is compilable. You can test this easily (and it should be a test case) with this modification to your code:

@jax.jit
def jax_normalize(to_transform):
  norm_jax, _, _ = normalizer_jax.normalize(to_transform)
  return norm_jax

norm_jax = jax_normalize(to_transform)

If that throws an error. there is more work to do. Usually, there will be an internal function you are calling that needs to be wrapped in a jax.jit, identifying what are the static_args.

If that does not throw an error, you might be done!

[swamidass]

So, running this, we do get an error. It's a classic example of where code needs to be refactored for JAX:

ODhat = OD[~jnp.any(OD < beta, axis=1)]

This needs to be refactored so that all intermediate matrices are a fixed size. The boolean select here yields an array with indeterminate size. The way to refactor this is with a mask.

I'm willing to make some of the changes, if I can. If i do, would you mind adding me as an author? I was also thinking of adding a utility we developed to enable parallelized application of this to large slides too.

[swamidass]

Also so a few of these

Inorm.at[Inorm > 255].set(255)

They are no ops as written. You'd need to do:

 Inorm = Inorm.at[Inorm > 255].set(255)

But I'm not srue they will compile, regardless. A better way to write that is:

Inorm = jnp.where(Inorm > 255, 255, Inorm) 

[andreped]

Its rather late in Norway, but I can try to incorporate your ideas, @swamidass. Give me a second, and I will make a commit if I get it working as intended.

[andreped]

Hmm, as far as I can tell, there does not seem to exist a masking mechanism in jax as of yet to mimic y = x[mask]. Rather strange I must say, but there are numerous threads about it.

Even more surprising: jax-ml/jax#11557

Any ideas? It is this line that is giving me a headache:

ODhat = OD[~jnp.any(OD < beta, axis=1)]

[andreped]

I got further, but the masking in Jax is giving me a headache. At least now, assuming there is a fix for this one line, adding the @jax.jit decorator works and hence, Macenko should work with Jax-backend.

You do not need to add it yourself externally, it is added directly within the class (see here).

Just reinstall the latest version of the same branch and run the same commands as mentioned previously to test: #31 (comment)

[andreped]

I'm willing to make some of the changes, if I can. If i do, would you mind adding me as an author? I was also thinking of adding a utility we developed to enable parallelized application of this to large slides too.

@swamidass Oh, and of course, contributors are always welcome. Regarding authorship, I am not the owner of this project, just contributing to it, but I'm open to the idea :] You have already been helpful in the Jax-backend implementation.

[swamidass]

I will take a look at it in a moment. but regarding:

ODhat = OD[~jnp.any(OD < beta, axis=1)]

The trick is to refactor the code so ODhat isn't needed. I believe the correct solution would be to change line 61,

_, eigvecs = jnp.linalg.eigh(jnp.cov(ODhat.T))

To, something close to:

mask = ~jnp.any(OD < beta, axis=1)
cov = jnp.cov(OD.T, fweights = mask)
_, eigvecs = jnp.linalg.eigh(cov)

That leaves leaves lines 27-32 to be refactored,

That = ODhat.dot(eigvecs[:, 1:3])

phi = jnp.arctan2(That[:, 1], That[:, 0])

minPhi = jnp.percentile(phi, alpha)
maxPhi = jnp.percentile(phi, 100 - alpha)

Into something like:

Th = OD.dot(eigvecs[:, 1:3])

phi = jnp.arctan2(Th[:, 1], Th[:, 0])

phi = jnp.where(mask, phi, jnp.inf)
pvalid = mask.mean() # proportion that is valid and not masked

minPhi = jnp.percentile(phi, alpha * pvalid)
maxPhi = jnp.percentile(phi, (100 - alpha) * pvalid)

I think those two changes will make the fix complete.