0

我对 Jax 文档感到困惑,这就是我想要做的:

def line(m,x,b):
  return m*x + b

grad(line)(1,2,3)

和错误:

---------------------------------------------------------------------------
FilteredStackTrace                        Traceback (most recent call last)
<ipython-input-48-d14b17620b30> in <module>()
      3 
----> 4 grad(line)(1,2,3)

FilteredStackTrace: TypeError: grad requires real- or complex-valued inputs (input dtype that is a sub-dtype of np.floating or np.complexfloating), but got int32. If you want to use integer-valued inputs, use vjp or set allow_int to True.

The stack trace above excludes JAX-internal frames.
The following is the original exception that occurred, unmodified.

--------------------

The above exception was the direct cause of the following exception:

TypeError                                 Traceback (most recent call last)
6 frames
/usr/local/lib/python3.7/dist-packages/jax/api.py in _check_input_dtype_revderiv(name, holomorphic, allow_int, x)
    844   elif not allow_int and not (dtypes.issubdtype(aval.dtype, np.floating) or
    845                               dtypes.issubdtype(aval.dtype, np.complexfloating)):
--> 846     raise TypeError(f"{name} requires real- or complex-valued inputs (input dtype that "
    847                     "is a sub-dtype of np.floating or np.complexfloating), "
    848                     f"but got {aval.dtype.name}. If you want to use integer-valued "

TypeError: grad requires real- or complex-valued inputs (input dtype that is a sub-dtype of np.floating or np.complexfloating), but got int32. If you want to use integer-valued inputs, use vjp or set allow_int to True.

我参考官方教程代码:

import jax.numpy as jnp
from jax import grad, jit, vmap
from jax import random

key = random.PRNGKey(0)

def sigmoid(x):
    return 0.5 * (jnp.tanh(x / 2) + 1)

# Outputs probability of a label being true.
def predict(W, b, inputs):
    return sigmoid(jnp.dot(inputs, W) + b)

# Build a toy dataset.
inputs = jnp.array([[0.52, 1.12,  0.77],
                   [0.88, -1.08, 0.15],
                   [0.52, 0.06, -1.30],
                   [0.74, -2.49, 1.39]])
targets = jnp.array([True, True, False, True])

# Training loss is the negative log-likelihood of the training examples.
def loss(W, b):
    preds = predict(W, b, inputs)
    label_probs = preds * targets + (1 - preds) * (1 - targets)
    return -jnp.sum(jnp.log(label_probs))

# Initialize random model coefficients
key, W_key, b_key = random.split(key, 3)
W = random.normal(W_key, (3,))
b = random.normal(b_key, ())

W_grad = grad(loss, argnums=0)(W, b)
print('W_grad', W_grad)

结果:

W_grad [-0.16965576 -0.8774648  -1.4901345 ]

我在这里做错了什么?我收集key到正在以某种重要的方式使用,但我无法弄清楚为什么/如何它是必要的。要回答这个问题,请根据需要调整第一个块中的代码以消除错误。

4

2 回答 2

2

Jax 告​​诉你它不喜欢整数。grad(line)(1.,2.,3.)(使用浮动)解决了这个问题。

于 2021-04-18T15:07:42.027 回答
2

我认为这里的错误很清楚:

TypeError: grad requires real- or complex-valued inputs (input dtype that is a sub-dtype of np.floating or np.complexfloating), but got int32. If you want to use integer-valued inputs, use vjp or set allow_int to True.

要使用grad(line)(1,2,3)with Int32,请将其更改为grad(line, allow_int=True)(1,2,3)

于 2021-04-18T15:12:49.847 回答