penzai: JAX research toolkit for building, editing, and visualizing neural nets

• penzai

  3 1,404 5.4 Python

  A JAX research toolkit for building, editing, and visualizing neural networks.

  

• Pytorch

  340 78,016 10.0 Python

  Tensors and Dynamic neural networks in Python with strong GPU acceleration

  

> does PyTorch have a similar concept

of course https://github.com/pytorch/pytorch/blob/main/torch/utils/_py...

• InfluxDB

  www.influxdata.com featured

  Power Real-Time Data Analytics at Scale. Get real-time insights from all types of time series data with InfluxDB. Ingest, query, and analyze billions of data points in real-time with unbounded cardinality.

  

• pytorch

  1 171 9.9 Python

  Ascend PyTorch adapter (torch_npu). Mirror of https://gitee.com/ascend/pytorch (by Ascend)

  

Its meteoric rise started well before the chip embargo. I've looked into it, it liberally borrows ideas from other frameworks, both PyTorch and Jax, and adds some of its own. You lose some of the conceptual purity, but it makes up for it in practical usability, assuming it works as it says on the tin, which it may or may not. PyTorch also has support for Ascend as far as I can tell https://github.com/Ascend/pytorch, so that support does not necessarily explain MindSpore's relative success. Why MindSpore is rising so rapidly is not entirely clear to me. Could be something as simple as preferring a domestic alternative that is adequate to the task and has better documentation in Chinese. Nowadays, however, I do agree that the various embargoes would help it (as well as Huawei) a great deal. As a side note I wish Huawei could export its silicon to the West. I bet that'd result in dramatically cheaper compute.

• pytreez

  1 1 - Python

  An implementation of Jax pytrees in pure python

I implemented Jax’s pytrees in pure python. You can use it with whatever you want. https://github.com/shawwn/pytreez

The readme is a todo, but the tests are complete. They’re the same that Jax itself uses, but zero dependencies. https://github.com/shawwn/pytreez/blob/master/tests/test_pyt...

The concept is simple. The hard part is cross pollination. Suppose you wanted to literally use Jax pytrees with PyTorch. Now you’ll have to import Jax, or my library, and register your modules with it. But anything else that ever uses pytrees need to use the same pytree library, because the registry (the thing that keeps track of pytree compatible classes) is in the library you choose. They don’t share registries.

A better way of phrasing it is that if you use a jax-style pytree interface, it should work with any other pytree library. But to my knowledge, the only pytree library besides Jax itself is mine here, and only I use it. So when you ask if pytree-compatible modules are compatible with PyTorch, it’s equivalent to asking whether PyTorch projects use jax, and the answer tends to be no.

Suggest a related project