DevOps研发效能
媒体矩阵
开源中国APP
授权协议 Apache-2.0
开发语言 Python 查看源码 »
操作系统 跨平台
软件类型 开源软件
开源组织 Google
地区 不详
投 递 者 h4cd
适用人群 未知
收录时间 2018-12-14

软件简介

JAX 是一个 TensorFlow 的简化库,它结合了 Autograd 和 XLA,专门用于高性能机器学习研究。

凭借 Autograd,JAX 可以求导循环、分支、递归和闭包函数,并且它可以进行三阶求导。通过 grad,它支持自动模式反向求导(反向传播)和正向求导,且二者可以任何顺序任意组合。

得力于 XLA,可以在 GPU 和 TPU 上编译和运行 NumPy 程序。默认情况下,编译发生在底层,库调用实时编译和执行。但是 JAX 还允许使用单一函数 API jit 将 Python 函数及时编译为 XLA 优化的内核。编译和自动求导可以任意组合,因此可以在 Python 环境下实现复杂的算法并获得最大的性能。

import jax.numpy as np
from jax import grad, jit, vmap
from functools import partial

def predict(params, inputs):
  for W, b in params:
    outputs = np.dot(inputs, W) + b
    inputs = np.tanh(outputs)
  return outputs

def logprob_fun(params, inputs, targets):
  preds = predict(params, inputs)
  return np.sum((preds - targets)**2)

grad_fun = jit(grad(logprob_fun))  # compiled gradient evaluation function
perex_grads = jit(vmap(grad_fun, in_axes=(None, 0, 0)))  # fast per-example grads

更深入地看,JAX 实际上是一个可扩展的可组合函数转换系统,grad 和 jit 都是这种转换的实例。

展开阅读全文

代码

的 Gitee 指数为
超过 的项目

评论

点击加入讨论🔥(1) 发布并加入讨论🔥
发表于行业趋势专区
2018/12/15 08:04

谷歌开源 TensorFlow 的简化库 JAX

谷歌开源了一个 TensorFlow 的简化库 JAX。 JAX 结合了 Autograd 和 XLA,专门用于高性能机器学习研究。 凭借 Autograd,JAX 可以求导循环、分支、递归和闭包函数,并且它可以进行三阶求导。通过 grad,它支持自动模式反向求导(反向传播)和正向求导,且二者可以任何顺序任意组合。 得力于 XLA,可以在 GPU 和 TPU 上编译和运行 NumPy 程序。默认情况下,编译发生在底层,库调用实时编译和执行。但是 JAX 还允许使用单一函数 ...

1
18
没有更多内容
加载失败,请刷新页面
点击加载更多
加载中
下一页
发表了博客
{{o.pubDate | formatDate}}

{{formatAllHtml(o.title)}}

{{parseInt(o.replyCount) | bigNumberTransform}}
{{parseInt(o.viewCount) | bigNumberTransform}}
没有更多内容
暂无内容
发表了问答
{{o.pubDate | formatDate}}

{{formatAllHtml(o.title)}}

{{parseInt(o.replyCount) | bigNumberTransform}}
{{parseInt(o.viewCount) | bigNumberTransform}}
没有更多内容
暂无内容
暂无内容
1 评论
11 收藏
分享
OSCHINA
登录后可查看更多优质内容
返回顶部
顶部