TensorFlow 简化库 JAX

TensorFlow 简化库 JAX

Apache-2.0
Python 查看源码»
跨平台
Google
2018-12-14
h4cd

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 都是这种转换的实例。

的码云指数为
超过 的项目
加载中

评论(1)

yl-yue
yl-yue
TensorFlow还未开始学习,这就出了新的了。😥

谷歌开源 TensorFlow 的简化库 JAX

谷歌开源了一个 TensorFlow 的简化库 JAX。 JAX 结合了 Autograd 和 XLA,专门用于高性能机器学习研究。 凭借 Autograd,JAX 可以求导循环、分支、递归和闭包函数,并且它可以进行三阶求导。...

2018/12/15 08:04

没有更多内容

加载失败,请刷新页面

没有更多内容

暂无问答

与Java Web Service相关的若干概念(JAX-WS,JAX-RS)

本篇关键词: JWS ,JAX-WS ,JAX-RS,REST,Restlet,SOAP 这一段时间从.NET转战java了。表示对java的一套相当不熟悉。新手,必须是新手。作为新手,我感到压力很大。 面对开发环境中的若干包...

2014/04/09 20:18
82
0
JAVA 三种WebService 规范

JAVA 中共有三种WebService 规范,分别是JAX-WS(JAX-RPC)、JAXM&SAAJ、JAX-RS。 1. Jaxws(掌握) JAX-WS 的全称为 Java API for XML-Based Webservices ,早期的基于SOAP 的JAVA 的Web 服务...

2018/07/21 19:17
53
0
java web service索引

正规军: 基于SOAP的JAX-RPC , JAX-WS, 基于RESTful的JAX-RS 杂牌:Apache系列。 书 RESTful Java with JAX-RS 2.0 RESTful Web Services with Dropwizard 以下这篇博客就是使用了JAX RS ht...

2014/06/11 00:48
121
0
JAX-WS和JAX-RS

JAX-WS  JAX-WS规范是一组XML web services的JAVA API,JAX-WS允许开发者可以选择RPC-oriented或者message-oriented 来实现自己的web services。   在 JAX-WS中,一个远程调用可以转换为一...

2013/11/28 16:17
382
0
java restful系列一 认识JAX-RS

1. 什么是JAX-RS JAX-RS是JAVA EE6 引入的一个新技术。 JAX-RS即Java API for RESTful Web Services,是一个Java 编程语言的应用程序接口,支持按照表述性状态转移(REST)架构风格创建Web服...

2016/11/02 23:51
12
0
A Comparison of Spring MVC and JAX-RS

Introduction Over the past few years REST has become an important concept influencing the design of web frameworks, web protocols and applications. If you haven't been exposed t...

2014/07/23 16:01
36
0

没有更多内容

加载失败,请刷新页面

返回顶部
顶部