Is there a path to compile and deploy JAX models on Android apps locally?
JAX – NumPy on the CPU, GPU, and TPU
61–70 of 147 posts
Re: JAX – NumPy on the CPU, GPU, and TPU
#62Is there a path to compile and deploy JAX models on Android apps locally?
I haven't tried it myself, but perhaps it's worth looking into the jax2tf -> tflite route.
Re: JAX – NumPy on the CPU, GPU, and TPU
#63Earlier quoted context omitted.
Having used JAX quite a bit for numerical computing (and having lectured on this use-case) I would say that a surprisingly large number of algorithms can be expressed as array[0] operations (even if it sometimes takes a bit of thinking). And, more importantly, things that cannot be expressed that way tend to not be a good fit for GPU computing anyway (independently of the language / framework you are using). [0]: `ar…
> things that cannot be expressed that way tend to not be a good fit for GPU computing anyway I'll have to disagree with you a little bit here. SIMT model of GPUs are quiet a bit more expressive than the numpy's SIMD model. As an obvious example, you'll have to manually maintain a mask to implement if/else i.e. code path divergence in SIMD. GPUs automatically does this and many more to make your life easier. And fran…
Convenient authoring doesn't necessarily make it a good fit for the hardware. Add in enough divergence and your GPU code is going to be matched or outperformed by a competent CPU implementation (on a chip of comparable size). Branchless code can result in substantial speedups on either.
To be fair though, modern GPUs are pretty good at branching and latency hiding, while numpy-style code has poor data locality unless you have a magic compiler.
Re: JAX – NumPy on the CPU, GPU, and TPU
#64I think JAX is cool, but I do find it slightly disingenuous when it claims to be "numpy by on the GPU" (as opposed to PyTorch), when actually there's a fundamental difference; it's functional. So if I have an array `x` and want to set index 0 to 10, I can't do: x[0] = 10 Instead I have to do: y = x.at[0].set(10) Of course this has advantages, but you can't then go and claim that JAX is a drop in replacement for numpy…
doesn't `[] =` just call a method on the object in python?
e.g, `x[0] = 10` is the same as `x.__set_item__(0, 10)`, so there shouldn't be any technical limitation to using `x[0]` (says the guy who never even imported jax)
Re: JAX – NumPy on the CPU, GPU, and TPU
#65I think JAX is cool, but I do find it slightly disingenuous when it claims to be "numpy by on the GPU" (as opposed to PyTorch), when actually there's a fundamental difference; it's functional. So if I have an array `x` and want to set index 0 to 10, I can't do: x[0] = 10 Instead I have to do: y = x.at[0].set(10) Of course this has advantages, but you can't then go and claim that JAX is a drop in replacement for numpy…
Sorry for my potentially VERY ignorant question, I only know functional programming at average joe level. Why can't you do the first in functional programming (not in this specific case because it's just how it is, but in general)? And even if you can't do so for any reasonable reason in functional (again, in general), what stops us to just add syntactic sugar to equal it to the second to make programmer's life easie…
There's 2 different aspects people mean when they call sth functional programming:
- higher order functions (lambdas, currying, closures, etc.)
- pure functions, immutability by default, side effects are pushed to the top level and marked clearly
The first aspect of functional programming has been already accepted by most OOP languages (even C++ has lambdas and closures).
The second aspect of functional programming is what makes it useful on GPU (because GPU architecture that makes it so powerful requires no interactions between code fragments that are run in parallel on 1000s of cores). So you can easily run pure functional code on GPU, but you can't easily run imperative code on GPU.
You can introduce side effects to functional programming, but then it ceases to be any more useful for GPU (and other parallel programming) than imperative/OOP.
Re: JAX – NumPy on the CPU, GPU, and TPU
#66I think JAX is cool, but I do find it slightly disingenuous when it claims to be "numpy by on the GPU" (as opposed to PyTorch), when actually there's a fundamental difference; it's functional. So if I have an array `x` and want to set index 0 to 10, I can't do: x[0] = 10 Instead I have to do: y = x.at[0].set(10) Of course this has advantages, but you can't then go and claim that JAX is a drop in replacement for numpy…
doesn't `[] =` just call a method on the object in python? e.g, `x[0] = 10` is the same as `x.__set_item__(0, 10)`, so there shouldn't be any technical limitation to using `x[0]` (says the guy who never even imported jax)
You could do `y = x.__setitem__(0, 10)`, but you cannot assign `x[0] = 10` to a new variable. If `__setitem__` was overridden, you would not be able to distinguish between these cases and raise an error in the second one.
Re: JAX – NumPy on the CPU, GPU, and TPU
#67Earlier quoted context omitted.
The fundamental reason why many functional languages won't allow you to do the first is that they use immutable data structures. We could indeed introduce syntactic sugar (`y= (x[0]:=10)` maybe), but you'll still need to introduce a new variable to hold the modified list.
It's my understanding that, at least in Python, you can't change immutable data type but you can just assign a new data to the same variable and therefore overwrite it, right? So even if JAX makes list type immutable, you can still just re-use `x` to save the new modified list.
[deleted]
Re: JAX – NumPy on the CPU, GPU, and TPU
#68> With its updated version of Autograd, JAX can automatically differentiate native Python and NumPy code.
Nice. When did they make this change?
Here is the old way in the docs, where you needed to define functions for the if-true branch and the if-false branch, and feed them to a conditional function, to get the normal if-then-else conditional.
https://jax.readthedocs.io/en/latest/notebooks/Common_Gotcha...
Re: JAX – NumPy on the CPU, GPU, and TPU
#69I think JAX is cool, but I do find it slightly disingenuous when it claims to be "numpy by on the GPU" (as opposed to PyTorch), when actually there's a fundamental difference; it's functional. So if I have an array `x` and want to set index 0 to 10, I can't do: x[0] = 10 Instead I have to do: y = x.at[0].set(10) Of course this has advantages, but you can't then go and claim that JAX is a drop in replacement for numpy…
On the other hand, if I wanted some scientific NumPy code to run on the GPU, I think rewriting it in JAX would probably be a better choice than PyTorch.
Re: JAX – NumPy on the CPU, GPU, and TPU
#70It took me a while to realize it, but Jax is actually a huge opportunity for a lot of scientific computing. Jax was originally developed as a more flexible platform for doing machine learning research. But Jax's real superpower is that it bundles XLA and makes it really easy to run computations on GPU or TPU. And huge swathes of scientific computation basically run large scale vectorized computations. When I was in a…
Would you happen to have sources on the three orders of magnitude speedup coming at no cost? I'd assume porting + data movement considerations making this task non-trivial.