6 ms·

It 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 pla

It 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 astronomy (about a decade ago) I did large scale simulations of gravitational interactions. But at the time all these simulations were done on CPU. Some of the really big efforts used more specialized chips, but it was a huge effort to write the code for it. But today with Jax, if you want to write an N-body simulation of a globular cluster, you can just code it up in numpy and it'll run on a GPU for free and be about 1000x faster. From what I can tell though, very few people in the sciences have caught on yet.

adgjlsfhk1 3y ago
the issue here is that if your ideal algorithm isn't simply expressible in numpy (which many aren't), you're pretty much out of luck. As a result, imo the better approach is to use a fast language that also compiles to GPU (e.g. Julia)
- krasin 3y ago
  An alternative is to write most of the program in Python + JAX + implement a few custom XLA ops in CUDA / Triton. That way, the program is very readable and can interoperate with the larger ecosystem, while still being fast to run.
- nestorD 3y ago
  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]: `array` is a shortcut here, JAX is not limited to operations on arrays.
  - sdenton4 3y ago
    Agreed. I've done a fair amount of reworking signal processing algorithms to run on GPU/TPU, and it's a different beast. You often have to really rebuild the algorithm from the ground up to take advantage of parallelization. But often you /can/ rework the algorithm, and end up with much higher throughput than the crusty old serial algorithm: there's typically nothing fundamentally stopping you from finding a good implementation, just that the original devs were working in the 70s and hasn't thought that far ahead.
  - shihab 3y ago
    > 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 frankly, I find it lot more easier to reason about what should happen to one data point than a bunch of them together. An interesting article I read recently that has some relevance to this discussion. https://pharr.org/matt/blog/2018/04/18/ispc-origins https://pharr.org/matt/blog/2018/04/18/ispc-origins
    - fulafel 3y ago
      To spell out what the linked ISPC post implies, most of the difference, like ISPC shows, is differences in GPU languages and compilers vs CPU side equivalents.
    - incrudible 3y ago
      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.
  - Certhas 3y ago
- marmaduke 3y ago
  Jax JIT of scan is fairly good, so loops aren't as slow as you'd expect.
  - cl3misch 3y ago
    A key difference is that each iteration of scan is called by the host. Put differently, JAX can't fuse scan into a single GPU kernel, but launches a kernel for each iteration. Depening on the workload this is no problem. If you have many cheap iterations, you will notice the overhead. I am not sure if they are working on fusing scan and what's the current status.
    - catgary 3y ago
      There is an ‘unroll’ parameter in scan that lets you control how many iterations of the loop are fused into a single kernel.
      - cl3misch 3y ago
        Yes, but is it really the same? Afaik the `unroll=n` parameter translates `n` iterations into a vanilla `for` loop which is then unrolled into sequential statements (in contrast to a JAX `fori` loop). There still is no loop on the accelerator, strictly speaking?
        marmaduke 3y ago
        I think this is up to XLA to handle not Jax. The whole selling point in TF of the tf.function decorator (which uses XLA underneath as well) is that it fuses arithmetic to lower launch count.
        marmaduke 3y ago
        there's a paper about it that I just found, enjoy https://arxiv.org/pdf/2301.13062 https://arxiv.org/pdf/2301.13062
tehsauce 3y ago
Jax is super useful for scientific computing. Although nbody sims might not be the best application. A naive nbody sim is very easy to implement and accelerate in jax (here’s my version: https://github.com/PWhiddy/jax-experiments/blob/main/nbody.ipynb https://github.com/PWhiddy/jax-experiments/blob/main/nbody.i...), but it can be tricky to scale it. This is because efficient nbody sims usually either rely on trees or spatial hashing/sorting which are tricky to efficiently implement with jax.
- mattjjatgoogle 3y ago
  Have you seen JAX MD? https://github.com/jax-md/jax-md https://github.com/jax-md/jax-md
  - tehsauce 3y ago
    I've seen it although haven't dived deep into it. It looks like they have some interesting support for particle cell data structures, but is fairly complicated and carries limitations: https://jax-md.readthedocs.io/en/main/_modules/jax_md/partition.html#cell_list https://jax-md.readthedocs.io/en/main/_modules/jax_md/partit...
  - dekhn 3y ago
    Last time I looked at JAX MD it didn't support most of the force field terms necessary for simulating proteins and DNA. For example, it could do n-body simulations with some potential, but not the bonds/torsions between atoms. It's unclear if they added support, but that's a huge gap in functionality compared to other systems.
dayeye2006 3y ago
Yes. But some of the algorithms cannot benefit that much from the GPU. In my field -- mathematical optimization, lots of algorithms rely on sparse matrix operations and takes many iterations until convergence.
- modeless 3y ago
  Would this help? https://jax.readthedocs.io/en/latest/jax.experimental.sparse.html https://jax.readthedocs.io/en/latest/jax.experimental.sparse...
  - marmaduke 3y ago
    Nope it's super slow for large sparse matrices. It's even faster to use generic scatter/gather to implement some, instead of that built in thing.
    - londons_explore 3y ago
      Have you investigated why? I know that many projects have an "implement first, optimize later" approach, and the lesser used functions might be far from optimal. Back in the tensorflow days, I had this issue and submitted a patch that gave a ~50x speedup for my usecase. It's always better to optimize the base function rather than have 100 people all manually working around the same performance issue.
      - marmaduke 3y ago
        Because they use a funny format (BCOO). I'm not mocking, it must be a solid choice for some reasons, like sparsification or other fancy stuff. But for large and even with batches (ie multiply with tall dense matrix), it doesn't match an equivalent scatter (x.at[idx].add(vals)). Which itself is several times slow than equivalent opencl (on an A40)
marmaduke 3y ago
None of these libraries unfortunately allows making good use of CPU vectorized units. Xla might produce some SIMD code but it pales in (performance) comparison to routines written explicitly for SIMD on GPU. ISPC is a good example of this.
swyx 3y ago
we did an interview with Chris Lattner of XLA fame where he also similarly had nice things to say about JAX: https://www.latent.space/p/modular https://www.latent.space/p/modular just sharing for those who want to learn more
- taktoa 3y ago
  The open source release of XLA predates Lattner's tenure at Google by 7 months, and it definitely existed before that -- the codebase was already 66k SLOC at that point. During his tenure it went from 100k SLOC to 250k SLOC. It's now 700k SLOC. He also has, as far as I can tell, zero commits in the XLA codebase. "Of LLVM fame" would be more accurate I think.
  - swyx 3y ago
    my bad - i guess i was just saying he led that team but didnt mean to imply he originated it you seem to have very precise knowledge of the SLOC at a point in time - just curious is there any tooling you used to do that? that can be pretty nifty to pull out on occasion
    - taktoa 3y ago
      I git cloned the repo and then ran sloccount after checking out various commits (just did `git log | grep -C3 'Jan 1 [0-9:]* 2017'` or similar to find the relevant commits)
      - swyx 3y ago
        ha, simple enough. thx
LightFog 3y ago
Porting existing codes is still a massive effort and there is low faith in long-term support from Google based software and hardware. I’m not aware on much (any) TPU use in scientific HPC.
- riedel 3y ago
  There is not yet. But there is huge pressure on HPC Centers at least in Europe to also make resources available for ML. Already many scientific supercomputers have GPUs as accelerators. So we might have the other situation: HPC users are faced more and more with machines with spare accelerators and it will make sense to use them. Actually it would totally make sense if JAX development is in part also public also financed in this case (e.g. through EuroHPC).
cl3misch 3y ago
Luckily I discovered JAX in the beginning of my PhD four years ago. It has made our data processing (biomedical imaging) so much easier and more readable, albeit with a slight learning curve due to JAX being functional/pure. I am also continously surprised how little adoption JIT and autodiff libraries have gotten in scientific computing. A lot of my colleagues somehow really like coding cost function gradients and fine-tuned GPU code by hand. I guess using something like JAX can reduce your standing in the group, because it can make it seem like coding algorithms is pretty easy.
- patrickkidger 3y ago
  I feel like I see the opposite -- that everything scientific computing is getting rewritten in something autodifferentiable! Whether that's JAX or something else. My experience might be biased though: shameless advert for Equinox (https://github.com/patrick-kidger/equinox https://github.com/patrick-kidger/equinox, 1.4k GitHub stars), which is now the foundation of quite a lot of SciComp in JAX. (Both internal and open-source.)
  - cl3misch 3y ago
    > foundation of quite a lot of SciComp in JAX ...if that SciComp uses machine learning, I guess? In my "physics of biomedical imaging" bubble, people are hardly doing state-of-the-art ML, but rather expensive forward models for which computing a gradient is cumbersome. But I know that e.g. Stephan Hoyer is a physicist and you are a mathematician originally -- I have read a lot of your JAX issues and libraries ;-) maybe it just depends on the "mini-bubble' aka. the indiviual research group and not only the field of science.
    - patrickkidger 3y ago
      > if that SciComp uses machine learning, I guess? Not necessarily! It's perfectly possible (and quite common) to e.g. write down a traditional parameterised ODE, and then optimise its parameters via gradient descent. Compute the gradients wrt parameters using autodiff through the numerical ODE solver. All without a single neural network in sight! ;) My usual spiel is that autodiff+autoparallel are really useful for any kind of numerical computation -- of which ML is a (popular, well funded) special case. At least in my mini bubble, these kinds of "scipy but autodifferentiable" use-cases are fairly common. > I have read a lot of your JAX issues and libraries ;-) Haha, that's fun to hear though! Thank you for sharing that.
__rito__ 3y ago
> "It took me a while to realize it, but Jax is actually a huge opportunity for a lot of scientific computing." In all conferences like NeurIPS, in Google ML Community days, etc., whenever there is a JAX workshop/tutorial/talk, it is always touted as a numerical computation library. And it was developed as such. Sure the focus is in ML, but everyone involved in it always have said that this is a general purpose scientific computing library. Flax, Haiku, etc. are Deep Learning libraries.
- matrss 3y ago
  Meanwhile, the first sentence in their readme is this: > JAX is Autograd and XLA, brought together for high-performance machine learning research. That does not really convey the generality of it that well.
  - mattjjatgoogle 3y ago
    You're right! Maybe we should revise that... I made https://github.com/google/jax/pull/17851 https://github.com/google/jax/pull/17851, comments welcome!
- londons_explore 3y ago
  So was tensorflow... And yet it's pretty much dead.
misja111 3y ago
> it'll run on a GPU for free and be about 1000x faster. Are there any benchmarks for that? Running on GPU never comes for free. You have to transfer data back and forth which has a cost, for instance.
- antognini 3y ago
  That was just from some quick benchmarks I did a few months back on some 10,000 particle N-body simulations. The performance boost will depend on the task, though. For the kinds of computations I did in grad school it would have been less effective since I was only looking at 3--5 objects, so there's just less parallelism to take advantage of.
konradha 3y ago
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.
dan_mctree 3y ago
It's so awkward that these truly fantastic tools for fast numerical computation (NumPy and JAX) have to be accessed through Python, which is a truly terrible language for fast numerical computation. Is anyone making any serious progress in fast GPU based computational tools for other faster languages? I'm looking for something that also works on the GPU on windows (unlike JAX)
- brahbrah 3y ago
  > have to be accessed through python It’s because most of the people doing these computations don’t have the capacity to become experts in multiple fields. They understand the math and analytics very well, and they expend all their time thinking about that, not about type systems, memory management, etc. Python lets them code without having to think about a lot of that stuff so they can focus on the things they care about. These aren’t computer scientists or programmers, they’re meteorologists, astronomers, oil and gas analysts, investment bankers etc. That’s why some truly great computer scientists and programmers invested their time into building these tools for python vs other languages.
- LaszloKv 3y ago
  You might want to look into Futhark https://futhark-lang.org/ https://futhark-lang.org/
- Q6T46nT668w6i3m 3y ago
  This type of code isn’t executed by the Python interpreter.
rcpt 3y ago
> large scale simulations of gravitational interactions I'm guessing this was mostly Fast Multipole Method? I don't think it ports that easily GPU since there's so much communication involved and the leaves don't do a whole lot
varelse 3y ago
[dead]
Q6T46nT668w6i3m 3y ago
Yep. I’ve been writing a molecular dynamics simulator in PyTorch and the reduction in scope is wild because of heterogenous operators and automatic differentiation.

adgjlsfhk1 3y ago

the issue here is that if your ideal algorithm isn't simply expressible in numpy (which many aren't), you're pretty much out of luck. As a result, imo the better approach is to use a fast language that also compiles to GPU (e.g. Julia)

krasin 3y ago
An alternative is to write most of the program in Python + JAX + implement a few custom XLA ops in CUDA / Triton. That way, the program is very readable and can interoperate with the larger ecosystem, while still being fast to run.
nestorD 3y ago
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]: `array` is a shortcut here, JAX is not limited to operations on arrays.
- sdenton4 3y ago
  Agreed. I've done a fair amount of reworking signal processing algorithms to run on GPU/TPU, and it's a different beast. You often have to really rebuild the algorithm from the ground up to take advantage of parallelization. But often you /can/ rework the algorithm, and end up with much higher throughput than the crusty old serial algorithm: there's typically nothing fundamentally stopping you from finding a good implementation, just that the original devs were working in the 70s and hasn't thought that far ahead.
- shihab 3y ago
  > 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 frankly, I find it lot more easier to reason about what should happen to one data point than a bunch of them together. An interesting article I read recently that has some relevance to this discussion. https://pharr.org/matt/blog/2018/04/18/ispc-origins https://pharr.org/matt/blog/2018/04/18/ispc-origins
  - fulafel 3y ago
    To spell out what the linked ISPC post implies, most of the difference, like ISPC shows, is differences in GPU languages and compilers vs CPU side equivalents.
  - incrudible 3y ago
    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.
- Certhas 3y ago
marmaduke 3y ago
Jax JIT of scan is fairly good, so loops aren't as slow as you'd expect.
- cl3misch 3y ago
  A key difference is that each iteration of scan is called by the host. Put differently, JAX can't fuse scan into a single GPU kernel, but launches a kernel for each iteration. Depening on the workload this is no problem. If you have many cheap iterations, you will notice the overhead. I am not sure if they are working on fusing scan and what's the current status.
  - catgary 3y ago
    There is an ‘unroll’ parameter in scan that lets you control how many iterations of the loop are fused into a single kernel.
    - cl3misch 3y ago
      Yes, but is it really the same? Afaik the `unroll=n` parameter translates `n` iterations into a vanilla `for` loop which is then unrolled into sequential statements (in contrast to a JAX `fori` loop). There still is no loop on the accelerator, strictly speaking?
      - marmaduke 3y ago
        I think this is up to XLA to handle not Jax. The whole selling point in TF of the tf.function decorator (which uses XLA underneath as well) is that it fuses arithmetic to lower launch count.
      - marmaduke 3y ago
        there's a paper about it that I just found, enjoy https://arxiv.org/pdf/2301.13062 https://arxiv.org/pdf/2301.13062

tehsauce 3y ago

Jax is super useful for scientific computing. Although nbody sims might not be the best application. A naive nbody sim is very easy to implement and accelerate in jax (here’s my version: https://github.com/PWhiddy/jax-experiments/blob/main/nbody.ipynb https://github.com/PWhiddy/jax-experiments/blob/main/nbody.i...), but it can be tricky to scale it. This is because efficient nbody sims usually either rely on trees or spatial hashing/sorting which are tricky to efficiently implement with jax.

mattjjatgoogle 3y ago
Have you seen JAX MD? https://github.com/jax-md/jax-md https://github.com/jax-md/jax-md
- tehsauce 3y ago
  I've seen it although haven't dived deep into it. It looks like they have some interesting support for particle cell data structures, but is fairly complicated and carries limitations: https://jax-md.readthedocs.io/en/main/_modules/jax_md/partition.html#cell_list https://jax-md.readthedocs.io/en/main/_modules/jax_md/partit...
- dekhn 3y ago
  Last time I looked at JAX MD it didn't support most of the force field terms necessary for simulating proteins and DNA. For example, it could do n-body simulations with some potential, but not the bonds/torsions between atoms. It's unclear if they added support, but that's a huge gap in functionality compared to other systems.

dayeye2006 3y ago

Yes. But some of the algorithms cannot benefit that much from the GPU. In my field -- mathematical optimization, lots of algorithms rely on sparse matrix operations and takes many iterations until convergence.

modeless 3y ago
Would this help? https://jax.readthedocs.io/en/latest/jax.experimental.sparse.html https://jax.readthedocs.io/en/latest/jax.experimental.sparse...
- marmaduke 3y ago
  Nope it's super slow for large sparse matrices. It's even faster to use generic scatter/gather to implement some, instead of that built in thing.
  - londons_explore 3y ago
    Have you investigated why? I know that many projects have an "implement first, optimize later" approach, and the lesser used functions might be far from optimal. Back in the tensorflow days, I had this issue and submitted a patch that gave a ~50x speedup for my usecase. It's always better to optimize the base function rather than have 100 people all manually working around the same performance issue.
    - marmaduke 3y ago
      Because they use a funny format (BCOO). I'm not mocking, it must be a solid choice for some reasons, like sparsification or other fancy stuff. But for large and even with batches (ie multiply with tall dense matrix), it doesn't match an equivalent scatter (x.at[idx].add(vals)). Which itself is several times slow than equivalent opencl (on an A40)

marmaduke 3y ago

None of these libraries unfortunately allows making good use of CPU vectorized units. Xla might produce some SIMD code but it pales in (performance) comparison to routines written explicitly for SIMD on GPU. ISPC is a good example of this.

swyx 3y ago

we did an interview with Chris Lattner of XLA fame where he also similarly had nice things to say about JAX: https://www.latent.space/p/modular https://www.latent.space/p/modular just sharing for those who want to learn more

taktoa 3y ago
The open source release of XLA predates Lattner's tenure at Google by 7 months, and it definitely existed before that -- the codebase was already 66k SLOC at that point. During his tenure it went from 100k SLOC to 250k SLOC. It's now 700k SLOC. He also has, as far as I can tell, zero commits in the XLA codebase. "Of LLVM fame" would be more accurate I think.
- swyx 3y ago
  my bad - i guess i was just saying he led that team but didnt mean to imply he originated it you seem to have very precise knowledge of the SLOC at a point in time - just curious is there any tooling you used to do that? that can be pretty nifty to pull out on occasion
  - taktoa 3y ago
    I git cloned the repo and then ran sloccount after checking out various commits (just did `git log | grep -C3 'Jan 1 [0-9:]* 2017'` or similar to find the relevant commits)
    - swyx 3y ago
      ha, simple enough. thx

LightFog 3y ago

Porting existing codes is still a massive effort and there is low faith in long-term support from Google based software and hardware. I’m not aware on much (any) TPU use in scientific HPC.

riedel 3y ago
There is not yet. But there is huge pressure on HPC Centers at least in Europe to also make resources available for ML. Already many scientific supercomputers have GPUs as accelerators. So we might have the other situation: HPC users are faced more and more with machines with spare accelerators and it will make sense to use them. Actually it would totally make sense if JAX development is in part also public also financed in this case (e.g. through EuroHPC).

cl3misch 3y ago

Luckily I discovered JAX in the beginning of my PhD four years ago. It has made our data processing (biomedical imaging) so much easier and more readable, albeit with a slight learning curve due to JAX being functional/pure. I am also continously surprised how little adoption JIT and autodiff libraries have gotten in scientific computing. A lot of my colleagues somehow really like coding cost function gradients and fine-tuned GPU code by hand. I guess using something like JAX can reduce your standing in the group, because it can make it seem like coding algorithms is pretty easy.

patrickkidger 3y ago
I feel like I see the opposite -- that everything scientific computing is getting rewritten in something autodifferentiable! Whether that's JAX or something else. My experience might be biased though: shameless advert for Equinox (https://github.com/patrick-kidger/equinox https://github.com/patrick-kidger/equinox, 1.4k GitHub stars), which is now the foundation of quite a lot of SciComp in JAX. (Both internal and open-source.)
- cl3misch 3y ago
  > foundation of quite a lot of SciComp in JAX ...if that SciComp uses machine learning, I guess? In my "physics of biomedical imaging" bubble, people are hardly doing state-of-the-art ML, but rather expensive forward models for which computing a gradient is cumbersome. But I know that e.g. Stephan Hoyer is a physicist and you are a mathematician originally -- I have read a lot of your JAX issues and libraries ;-) maybe it just depends on the "mini-bubble' aka. the indiviual research group and not only the field of science.
  - patrickkidger 3y ago
    > if that SciComp uses machine learning, I guess? Not necessarily! It's perfectly possible (and quite common) to e.g. write down a traditional parameterised ODE, and then optimise its parameters via gradient descent. Compute the gradients wrt parameters using autodiff through the numerical ODE solver. All without a single neural network in sight! ;) My usual spiel is that autodiff+autoparallel are really useful for any kind of numerical computation -- of which ML is a (popular, well funded) special case. At least in my mini bubble, these kinds of "scipy but autodifferentiable" use-cases are fairly common. > I have read a lot of your JAX issues and libraries ;-) Haha, that's fun to hear though! Thank you for sharing that.

__rito__ 3y ago

> "It took me a while to realize it, but Jax is actually a huge opportunity for a lot of scientific computing." In all conferences like NeurIPS, in Google ML Community days, etc., whenever there is a JAX workshop/tutorial/talk, it is always touted as a numerical computation library. And it was developed as such. Sure the focus is in ML, but everyone involved in it always have said that this is a general purpose scientific computing library. Flax, Haiku, etc. are Deep Learning libraries.

matrss 3y ago
Meanwhile, the first sentence in their readme is this: > JAX is Autograd and XLA, brought together for high-performance machine learning research. That does not really convey the generality of it that well.
- mattjjatgoogle 3y ago
  You're right! Maybe we should revise that... I made https://github.com/google/jax/pull/17851 https://github.com/google/jax/pull/17851, comments welcome!
londons_explore 3y ago
So was tensorflow... And yet it's pretty much dead.

misja111 3y ago

> it'll run on a GPU for free and be about 1000x faster. Are there any benchmarks for that? Running on GPU never comes for free. You have to transfer data back and forth which has a cost, for instance.

antognini 3y ago
That was just from some quick benchmarks I did a few months back on some 10,000 particle N-body simulations. The performance boost will depend on the task, though. For the kinds of computations I did in grad school it would have been less effective since I was only looking at 3--5 objects, so there's just less parallelism to take advantage of.

konradha 3y ago

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.

dan_mctree 3y ago

It's so awkward that these truly fantastic tools for fast numerical computation (NumPy and JAX) have to be accessed through Python, which is a truly terrible language for fast numerical computation. Is anyone making any serious progress in fast GPU based computational tools for other faster languages? I'm looking for something that also works on the GPU on windows (unlike JAX)

brahbrah 3y ago
> have to be accessed through python It’s because most of the people doing these computations don’t have the capacity to become experts in multiple fields. They understand the math and analytics very well, and they expend all their time thinking about that, not about type systems, memory management, etc. Python lets them code without having to think about a lot of that stuff so they can focus on the things they care about. These aren’t computer scientists or programmers, they’re meteorologists, astronomers, oil and gas analysts, investment bankers etc. That’s why some truly great computer scientists and programmers invested their time into building these tools for python vs other languages.
LaszloKv 3y ago
You might want to look into Futhark https://futhark-lang.org/ https://futhark-lang.org/
Q6T46nT668w6i3m 3y ago
This type of code isn’t executed by the Python interpreter.

rcpt 3y ago

> large scale simulations of gravitational interactions I'm guessing this was mostly Fast Multipole Method? I don't think it ports that easily GPU since there's so much communication involved and the leaves don't do a whole lot

varelse 3y ago

[dead]

Q6T46nT668w6i3m 3y ago

Yep. I’ve been writing a molecular dynamics simulator in PyTorch and the reduction in scope is wild because of heterogenous operators and automatic differentiation.