这个 Jax 函数的纯函数版本是什么？

问：

我正在编写一个简单的分类器，试图掌握 Jax。我想对数据进行归一化（它是来自 sklearn 的 Iris 数据集）并且我的小函数可以工作，但从我在 Jax the Sharp Bits 文档中读到的内容来看，我应该避免使用 lambda 函数和迭代向量。我不精通函数式编程，我很好奇是否有更好、更像 Jax 的方法来做到这一点。这是我到目前为止的代码：

import jax.numpy as jnp 
from jax import jit, vmap 
# lots of imports ... 

iris = load_it('data', 'iris.pkl')

def normalize(data): 
    return jnp.apply_along_axis(lambda x: x/jnp.linalg.norm(x), 1, data) 

# TODO: use a functional style, maybe use partial
# and get rid of the lambda ...

tic = time.perf_counter() 
iris_data_normal = normalize(iris.data) 
toc = time.perf_counter() 
print(f"It took jax {toc - tic:0.4f} seconds.")

当我运行它时，我得到：任何指导都是最值得赞赏的！It took jax 0.0677 seconds.