科学编程系列之 01 在对PyTorch优化 Python 脚本感到失望后，我开始使用JAX进行分子动力学展示，高性能数值计算库，提供自动微分和在 CPU 和 GPU 上的优化执行

简介

在对PyTorch优化 Python 脚本感到失望后，我开始使用JAX。我的项目涉及两个主要部分：根据原子位置计算描述符，并使用这些描述符作为多个神经网络的输入，以预测粒子系统的总势能和力。虽然神经网络部分足够快，但即使使用TorchScript后，描述符计算（尤其是梯度评估）也无法高效执行。在寻找支持自动微分的 Python 替代框架时，我发现了 JAX。事实证明，它对于构建（物理感知）机器学习模型非常有效，提供了我所需的灵活性和性能。

推荐文章

• 《AnythingLLM教程系列之 05 AnythingLLM 允许您创建嵌入式聊天小部件，可以使用简单的标签轻松集成到任何网站中》 权重1，本地类、AnythingLLM

• 《大模型系统之 01 duckdb-nsql-7B 的 7B 参数文本到 :Text2SQL OpenSource：duckdb-nsql-7B，在本地设置上带有 Ollama 和 LlamaIndex》 权重1，本地类、duckdb-nsql-7B、Text2SQL、Ollama、LlamaIndex