Skip to content

P
pytensor
提交 ffb72e42 authored 作者: Virgile Andreani's avatar Virgile Andreani 提交者： Thomas Wiecki
浏览文件
操作

Enable ruff to format code in docstrings

上级 86c8a00f
显示空白字符变更
内嵌 并排
正在显示 包含 192 行增加 和 123 行删除
pyproject.toml
浏览文件 @ ffb72e42
......@@ -124,6 +124,9 @@ testpaths = "tests/"
line-length = 88
exclude = ["doc/", "pytensor/_version.py"]
[tool.ruff.format]
docstring-code-format = true
[tool.ruff.lint]
select = ["C", "E", "F", "I", "UP", "W", "RUF", "PERF", "PTH", "ISC"]
ignore = ["C408", "C901", "E501", "E741", "RUF012", "PERF203", "ISC001"]
......
pytensor/compile/builders.py
浏览文件 @ ffb72e42
......@@ -190,7 +190,8 @@ class OpFromGraph(Op, HasInnerGraph):
from pytensor import function, tensor as pt
from pytensor.compile.builders import OpFromGraph
x, y, z = pt.scalars('xyz')
x, y, z = pt.scalars("xyz")
e = x + y * z
op = OpFromGraph([x, y, z], [e])
# op behaves like a normal pytensor op
......@@ -206,7 +207,7 @@ class OpFromGraph(Op, HasInnerGraph):
from pytensor import config, function, tensor as pt
from pytensor.compile.builders import OpFromGraph
x, y, z = pt.scalars('xyz')
x, y, z = pt.scalars("xyz")
s = pytensor.shared(np.random.random((2, 2)).astype(config.floatX))
e = x + y * z + s
op = OpFromGraph([x, y, z], [e])
......@@ -221,12 +222,16 @@ class OpFromGraph(Op, HasInnerGraph):
from pytensor import function, tensor as pt, grad
from pytensor.compile.builders import OpFromGraph
x, y, z = pt.scalars('xyz')
x, y, z = pt.scalars("xyz")
e = x + y * z
def rescale_dy(inps, outputs, out_grads):
x, y, z = inps
g, = out_grads
return z*2
(g,) = out_grads
return z * 2
op = OpFromGraph(
[x, y, z],
[e],
......@@ -236,7 +241,7 @@ class OpFromGraph(Op, HasInnerGraph):
dx, dy, dz = grad(e2, [x, y, z])
fn = function([x, y, z], [dx, dy, dz])
# the gradient wrt y is now doubled
fn(2., 3., 4.) # [1., 8., 3.]
fn(2.0, 3.0, 4.0) # [1., 8., 3.]
"""
......
pytensor/gradient.py
浏览文件 @ ffb72e42
......@@ -692,25 +692,24 @@ def subgraph_grad(wrt, end, start=None, cost=None, details=False):
.. code-block:: python
x, t = pytensor.tensor.fvector('x'), pytensor.tensor.fvector('t')
w1 = pytensor.shared(np.random.standard_normal((3,4)))
w2 = pytensor.shared(np.random.standard_normal((4,2)))
a1 = pytensor.tensor.tanh(pytensor.tensor.dot(x,w1))
a2 = pytensor.tensor.tanh(pytensor.tensor.dot(a1,w2))
x, t = pytensor.tensor.fvector("x"), pytensor.tensor.fvector("t")
w1 = pytensor.shared(np.random.standard_normal((3, 4)))
w2 = pytensor.shared(np.random.standard_normal((4, 2)))
a1 = pytensor.tensor.tanh(pytensor.tensor.dot(x, w1))
a2 = pytensor.tensor.tanh(pytensor.tensor.dot(a1, w2))
cost2 = pytensor.tensor.sqr(a2 - t).sum()
cost2 += pytensor.tensor.sqr(w2.sum())
cost1 = pytensor.tensor.sqr(w1.sum())
params = [[w2],[w1]]
costs = [cost2,cost1]
params = [[w2], [w1]]
costs = [cost2, cost1]
grad_ends = [[a1], [x]]
next_grad = None
param_grads = []
for i in range(2):
param_grad, next_grad = pytensor.subgraph_grad(
wrt=params[i], end=grad_ends[i],
start=next_grad, cost=costs[i]
wrt=params[i], end=grad_ends[i], start=next_grad, cost=costs[i]
)
next_grad = dict(zip(grad_ends[i], next_grad))
param_grads.extend(param_grad)
......@@ -1704,9 +1703,11 @@ def verify_grad(
Examples
--------
>>> verify_grad(pytensor.tensor.tanh,
>>> verify_grad(
... pytensor.tensor.tanh,
... (np.asarray([[2, 3, 4], [-1, 3.3, 9.9]]),),
... rng=np.random.default_rng(23098))
... rng=np.random.default_rng(23098),
... )
Parameters
----------
......@@ -2342,9 +2343,9 @@ def grad_clip(x, lower_bound, upper_bound):
Examples
--------
>>> x = pytensor.tensor.type.scalar()
>>> z = pytensor.gradient.grad(grad_clip(x, -1, 1)**2, x)
>>> z = pytensor.gradient.grad(grad_clip(x, -1, 1) ** 2, x)
>>> z2 = pytensor.gradient.grad(x**2, x)
>>> f = pytensor.function([x], outputs = [z, z2])
>>> f = pytensor.function([x], outputs=[z, z2])
>>> print(f(2.0))
[array(1.), array(4.)]
......@@ -2383,7 +2384,7 @@ def grad_scale(x, multiplier):
>>> fprime = pytensor.function([x], fp)
>>> print(fprime(2)) # doctest: +ELLIPSIS
-0.416...
>>> f_inverse=grad_scale(fx, -1.)
>>> f_inverse = grad_scale(fx, -1.0)
>>> fpp = pytensor.grad(f_inverse, wrt=x)
>>> fpprime = pytensor.function([x], fpp)
>>> print(fpprime(2)) # doctest: +ELLIPSIS
......
pytensor/graph/basic.py
浏览文件 @ ffb72e42
......@@ -404,13 +404,19 @@ class Variable(Node, Generic[_TypeType, OptionalApplyType]):
c = a + b # create a simple expression
f = pytensor.function([b], [c]) # this works because a has a value associated with it already
f = pytensor.function(
[b], [c]
) # this works because a has a value associated with it already
assert 4.0 == f(2.5) # bind 2.5 to an internal copy of b and evaluate an internal c
pytensor.function([a], [c]) # compilation error because b (required by c) is undefined
pytensor.function(
[a], [c]
) # compilation error because b (required by c) is undefined
pytensor.function([a,b], [c]) # compilation error because a is constant, it can't be an input
pytensor.function(
[a, b], [c]
) # compilation error because a is constant, it can't be an input
The python variables ``a, b, c`` all refer to instances of type
......@@ -587,10 +593,10 @@ class Variable(Node, Generic[_TypeType, OptionalApplyType]):
>>> import numpy as np
>>> import pytensor.tensor as pt
>>> x = pt.dscalar('x')
>>> y = pt.dscalar('y')
>>> x = pt.dscalar("x")
>>> y = pt.dscalar("y")
>>> z = x + y
>>> np.allclose(z.eval({x : 16.3, y : 12.1}), 28.4)
>>> np.allclose(z.eval({x: 16.3, y: 12.1}), 28.4)
True
We passed :meth:`eval` a dictionary mapping symbolic PyTensor
......@@ -963,9 +969,9 @@ def explicit_graph_inputs(
import pytensor.tensor as pt
from pytensor.graph.basic import explicit_graph_inputs
x = pt.vector('x')
x = pt.vector("x")
y = pt.constant(2)
z = pt.mul(x*y)
z = pt.mul(x * y)
inputs = list(explicit_graph_inputs(z))
f = pytensor.function(inputs, z)
......@@ -1041,7 +1047,7 @@ def orphans_between(
>>> from pytensor.graph.basic import orphans_between
>>> from pytensor.tensor import scalars
>>> x, y = scalars("xy")
>>> list(orphans_between([x], [(x+y)]))
>>> list(orphans_between([x], [(x + y)]))
[y]
"""
......
pytensor/link/c/interface.py
浏览文件 @ ffb72e42
......@@ -30,7 +30,7 @@ class CLinkerObject:
.. code-block:: python
def c_headers(self, **kwargs):
return ['<iostream>', '<math.h>', '/full/path/to/header.h']
return ["<iostream>", "<math.h>", "/full/path/to/header.h"]
"""
......@@ -54,7 +54,7 @@ class CLinkerObject:
.. code-block:: python
def c_header_dirs(self, **kwargs):
return ['/usr/local/include', '/opt/weirdpath/src/include']
return ["/usr/local/include", "/opt/weirdpath/src/include"]
"""
return []
......@@ -134,7 +134,7 @@ class CLinkerObject:
.. code-block:: python
def c_compile_args(self, **kwargs):
return ['-ffast-math']
return ["-ffast-math"]
"""
return []
......
pytensor/link/c/params_type.py
浏览文件 @ ffb72e42
......@@ -29,7 +29,9 @@ In your Op sub-class:
.. code-block:: python
params_type = ParamsType(attr1=TensorType('int32', shape=(None, None)), attr2=ScalarType('float64'))
params_type = ParamsType(
attr1=TensorType("int32", shape=(None, None)), attr2=ScalarType("float64")
)
If your op contains attributes ``attr1`` **and** ``attr2``, the default ``op.get_params()``
implementation will automatically try to look for it and generate an appropriate Params object.
......@@ -77,26 +79,35 @@ enumerations will be directly available as ParamsType attributes.
from pytensor.link.c.params_type import ParamsType
from pytensor.link.c.type import EnumType, EnumList
wrapper = ParamsType(enum1=EnumList('CONSTANT_1', 'CONSTANT_2', 'CONSTANT_3'),
enum2=EnumType(PI=3.14, EPSILON=0.001))
wrapper = ParamsType(
enum1=EnumList("CONSTANT_1", "CONSTANT_2", "CONSTANT_3"),
enum2=EnumType(PI=3.14, EPSILON=0.001),
)
# Each enum constant is available as a wrapper attribute:
print(wrapper.CONSTANT_1, wrapper.CONSTANT_2, wrapper.CONSTANT_3,
wrapper.PI, wrapper.EPSILON)
print(
wrapper.CONSTANT_1,
wrapper.CONSTANT_2,
wrapper.CONSTANT_3,
wrapper.PI,
wrapper.EPSILON,
)
# For convenience, you can also look for a constant by name with
# ``ParamsType.get_enum()`` method.
pi = wrapper.get_enum('PI')
epsilon = wrapper.get_enum('EPSILON')
constant_2 = wrapper.get_enum('CONSTANT_2')
pi = wrapper.get_enum("PI")
epsilon = wrapper.get_enum("EPSILON")
constant_2 = wrapper.get_enum("CONSTANT_2")
print(pi, epsilon, constant_2)
This implies that a ParamsType cannot contain different enum types with common enum names::
# Following line will raise an error,
# as there is a "CONSTANT_1" defined both in enum1 and enum2.
wrapper = ParamsType(enum1=EnumList('CONSTANT_1', 'CONSTANT_2'),
enum2=EnumType(CONSTANT_1=0, CONSTANT_3=5))
wrapper = ParamsType(
enum1=EnumList("CONSTANT_1", "CONSTANT_2"),
enum2=EnumType(CONSTANT_1=0, CONSTANT_3=5),
)
If your enum types contain constant aliases, you can retrieve them from ParamsType
with ``ParamsType.enum_from_alias(alias)`` method (see :class:`pytensor.link.c.type.EnumType`
......@@ -104,11 +115,12 @@ for more info about enumeration aliases).
.. code-block:: python
wrapper = ParamsType(enum1=EnumList('A', ('B', 'beta'), 'C'),
enum2=EnumList(('D', 'delta'), 'E', 'F'))
wrapper = ParamsType(
enum1=EnumList("A", ("B", "beta"), "C"), enum2=EnumList(("D", "delta"), "E", "F")
)
b1 = wrapper.B
b2 = wrapper.get_enum('B')
b3 = wrapper.enum_from_alias('beta')
b2 = wrapper.get_enum("B")
b3 = wrapper.enum_from_alias("beta")
assert b1 == b2 == b3
"""
......@@ -236,10 +248,13 @@ class Params(dict):
from pytensor.link.c.params_type import ParamsType, Params
from pytensor.scalar import ScalarType
# You must create a ParamsType first:
params_type = ParamsType(attr1=ScalarType('int32'),
key2=ScalarType('float32'),
field3=ScalarType('int64'))
params_type = ParamsType(
attr1=ScalarType("int32"),
key2=ScalarType("float32"),
field3=ScalarType("int64"),
)
# Then you can create a Params object with
# the params type defined above and values for attributes.
params = Params(params_type, attr1=1, key2=2.0, field3=3)
......@@ -491,11 +506,13 @@ class ParamsType(CType):
from pytensor.link.c.type import EnumType, EnumList
from pytensor.scalar import ScalarType
wrapper = ParamsType(scalar=ScalarType('int32'),
wrapper = ParamsType(
scalar=ScalarType("int32"),
letters=EnumType(A=1, B=2, C=3),
digits=EnumList('ZERO', 'ONE', 'TWO'))
print(wrapper.get_enum('C')) # 3
print(wrapper.get_enum('TWO')) # 2
digits=EnumList("ZERO", "ONE", "TWO"),
)
print(wrapper.get_enum("C")) # 3
print(wrapper.get_enum("TWO")) # 2
# You can also directly do:
print(wrapper.C)
......@@ -520,17 +537,19 @@ class ParamsType(CType):
from pytensor.link.c.type import EnumType, EnumList
from pytensor.scalar import ScalarType
wrapper = ParamsType(scalar=ScalarType('int32'),
letters=EnumType(A=(1, 'alpha'), B=(2, 'beta'), C=3),
digits=EnumList(('ZERO', 'nothing'), ('ONE', 'unit'), ('TWO', 'couple')))
print(wrapper.get_enum('C')) # 3
print(wrapper.get_enum('TWO')) # 2
print(wrapper.enum_from_alias('alpha')) # 1
print(wrapper.enum_from_alias('nothing')) # 0
wrapper = ParamsType(
scalar=ScalarType("int32"),
letters=EnumType(A=(1, "alpha"), B=(2, "beta"), C=3),
digits=EnumList(("ZERO", "nothing"), ("ONE", "unit"), ("TWO", "couple")),
)
print(wrapper.get_enum("C")) # 3
print(wrapper.get_enum("TWO")) # 2
print(wrapper.enum_from_alias("alpha")) # 1
print(wrapper.enum_from_alias("nothing")) # 0
# For the following, alias 'C' is not defined, so the method looks for
# a constant named 'C', and finds it.
print(wrapper.enum_from_alias('C')) # 3
print(wrapper.enum_from_alias("C")) # 3
.. note::
......@@ -567,12 +586,14 @@ class ParamsType(CType):
from pytensor.tensor.type import dmatrix
from pytensor.scalar import ScalarType
class MyObject:
def __init__(self):
self.a = 10
self.b = numpy.asarray([[1, 2, 3], [4, 5, 6]])
params_type = ParamsType(a=ScalarType('int32'), b=dmatrix, c=ScalarType('bool'))
params_type = ParamsType(a=ScalarType("int32"), b=dmatrix, c=ScalarType("bool"))
o = MyObject()
value_for_c = False
......
pytensor/link/c/type.py
浏览文件 @ ffb72e42
......@@ -318,7 +318,7 @@ class EnumType(CType, dict):
.. code-block:: python
enum = EnumType(CONSTANT_1=1, CONSTANT_2=2.5, CONSTANT_3=False, CONSTANT_4=True)
print (enum.CONSTANT_1, enum.CONSTANT_2, enum.CONSTANT_3, enum.CONSTANT_4)
print(enum.CONSTANT_1, enum.CONSTANT_2, enum.CONSTANT_3, enum.CONSTANT_4)
# will print 1 2.5 0 1
In C code:
......@@ -334,7 +334,7 @@ class EnumType(CType, dict):
.. code-block:: python
enum = EnumType(CONSTANT_1=0, CONSTANT_2=1, CONSTANT_3=2, ctype='size_t')
enum = EnumType(CONSTANT_1=0, CONSTANT_2=1, CONSTANT_3=2, ctype="size_t")
# In C code, the Op param will then be a ``size_t``.
.. note::
......@@ -349,8 +349,9 @@ class EnumType(CType, dict):
.. code-block:: python
enum = EnumType(CONSTANT_1=0, CONSTANT_2=1, CONSTANT_3=2,
ctype='size_t', cname='MyEnumName')
enum = EnumType(
CONSTANT_1=0, CONSTANT_2=1, CONSTANT_3=2, ctype="size_t", cname="MyEnumName"
)
**Example with aliases**
......@@ -359,7 +360,7 @@ class EnumType(CType, dict):
To give an alias to a constant in the EnumType constructor, use the following key-value syntax::
constant_name=(constant_alias, constant_value)
constant_name = (constant_alias, constant_value)
You can then retrieve a constant from an alias with method ``EnumType.fromalias()``.
......@@ -372,23 +373,23 @@ class EnumType(CType, dict):
from pytensor.link.c.type import EnumType
# You can remark that constant 'C' does not have an alias.
enum = EnumType(A=('alpha', 1), B=('beta', 2), C=3, D=('delta', 4))
enum = EnumType(A=("alpha", 1), B=("beta", 2), C=3, D=("delta", 4))
# Constants are all directly available by name.
print(enum.A, enum.B, enum.C, enum.D)
# But we can also now get some constants by alias.
a = enum.fromalias('alpha')
b = enum.fromalias('beta')
d = enum.fromalias('delta')
a = enum.fromalias("alpha")
b = enum.fromalias("beta")
d = enum.fromalias("delta")
# If method fromalias() receives an unknown alias,
# it will looks for a constant with this alias
# as exact constant name.
c = enum.fromalias('C') # will get enum.C
c = enum.fromalias("C") # will get enum.C
# An alias defined in an EnumType will be correctly converted with non-strict filtering.
value = enum.filter('delta', strict=False)
value = enum.filter("delta", strict=False)
# value now contains enum.D, ie. 4.
.. note::
......@@ -648,14 +649,24 @@ class EnumList(EnumType):
Example::
enum = EnumList('CONSTANT_1', 'CONSTANT_2', 'CONSTANT_3', 'CONSTANT_4', 'CONSTANT_5')
print (enum.CONSTANT_1, enum.CONSTANT_2, enum.CONSTANT_3, enum.CONSTANT_4, enum.CONSTANT_5)
enum = EnumList(
"CONSTANT_1", "CONSTANT_2", "CONSTANT_3", "CONSTANT_4", "CONSTANT_5"
)
print(
enum.CONSTANT_1,
enum.CONSTANT_2,
enum.CONSTANT_3,
enum.CONSTANT_4,
enum.CONSTANT_5,
)
# will print: 0 1 2 3 4
Like :class:`EnumType`, you can also define the C type and a C name for the op param.
Default C type is ``int``::
enum = EnumList('CONSTANT_1', 'CONSTANT_2', 'CONSTANT_3', 'CONSTANT_4', ctype='unsigned int')
enum = EnumList(
"CONSTANT_1", "CONSTANT_2", "CONSTANT_3", "CONSTANT_4", ctype="unsigned int"
)
Like :class:`EnumType`, you can also add an alias to a constant, by replacing the only constant name
(e.g. ``'CONSTANT_NAME'``) by a couple with constant name first and constant alias second
......@@ -663,7 +674,7 @@ class EnumList(EnumType):
.. code-block:: python
enum = EnumList(('A', 'alpha'), ('B', 'beta'), 'C', 'D', 'E', 'F', ('G', 'gamma'))
enum = EnumList(("A", "alpha"), ("B", "beta"), "C", "D", "E", "F", ("G", "gamma"))
See test class :class:`tests.graph.test_types.TestOpEnumList` for a working example.
......@@ -727,7 +738,9 @@ class CEnumType(EnumList):
.. code-block:: python
enum = CEnumType('CONSTANT_CNAME_1', 'CONSTANT_CNAME_2', 'CONSTANT_CNAME_3', ctype='long')
enum = CEnumType(
"CONSTANT_CNAME_1", "CONSTANT_CNAME_2", "CONSTANT_CNAME_3", ctype="long"
)
Like :class:`EnumList`, you can also add an alias to a constant, with same syntax as in :class:`EnumList`.
......
pytensor/link/numba/dispatch/elemwise.py
浏览文件 @ ffb72e42
......@@ -185,7 +185,9 @@ def create_axis_reducer(
.. code-block:: python
def careduce_axis(x):
res_shape = tuple(shape[i] if i < axis else shape[i + 1] for i in range(ndim - 1))
res_shape = tuple(
shape[i] if i < axis else shape[i + 1] for i in range(ndim - 1)
)
res = np.full(res_shape, identity, dtype=dtype)
x_axis_first = x.transpose(reaxis_first)
......
pytensor/printing.py
浏览文件 @ ffb72e42
......@@ -1247,10 +1247,11 @@ def pydotprint(
.. code-block:: python
import pytensor
v = pytensor.tensor.vector()
from IPython.display import SVG
SVG(pytensor.printing.pydotprint(v*2, return_image=True,
format='svg'))
SVG(pytensor.printing.pydotprint(v * 2, return_image=True, format="svg"))
In the graph, ellipses are Apply Nodes (the execution of an op)
and boxes are variables. If variables have names they are used as
......
pytensor/scalar/basic.py
浏览文件 @ ffb72e42
......@@ -209,9 +209,9 @@ class autocast_float_as:
Examples
--------
>>> from pytensor.tensor import fvector
>>> with autocast_float_as('float32'):
... assert (fvector() + 1.1).dtype == 'float32' # temporary downcasting
>>> assert (fvector() + 1.1).dtype == 'float64' # back to default behaviour
>>> with autocast_float_as("float32"):
... assert (fvector() + 1.1).dtype == "float32" # temporary downcasting
>>> assert (fvector() + 1.1).dtype == "float64" # back to default behaviour
"""
......
pytensor/scan/basic.py
浏览文件 @ ffb72e42
......@@ -207,13 +207,20 @@ def scan(
.. code-block:: python
scan(fn, sequences = [ dict(input= Sequence1, taps = [-3,2,-1])
, Sequence2
, dict(input = Sequence3, taps = 3) ]
, outputs_info = [ dict(initial = Output1, taps = [-3,-5])
, dict(initial = Output2, taps = None)
, Output3 ]
, non_sequences = [ Argument1, Argument2])
scan(
fn,
sequences=[
dict(input=Sequence1, taps=[-3, 2, -1]),
Sequence2,
dict(input=Sequence3, taps=3),
],
outputs_info=[
dict(initial=Output1, taps=[-3, -5]),
dict(initial=Output2, taps=None),
Output3,
],
non_sequences=[Argument1, Argument2],
)
`fn` should expect the following arguments in this given order:
......@@ -243,8 +250,9 @@ def scan(
W = pt.matrix()
W_2 = W**2
def f(x):
return pt.dot(x,W_2)
return pt.dot(x, W_2)
The function `fn` is expected to return two things. One is a list of
outputs ordered in the same order as `outputs_info`, with the
......@@ -266,7 +274,7 @@ def scan(
.. code-block:: python
...
return [y1_t, y2_t], {x:x+1}, until(x < 50)
return [y1_t, y2_t], {x: x + 1}, until(x < 50)
Note that a number of steps--considered in here as the maximum
number of steps--is still required even though a condition is
......
pytensor/tensor/basic.py
浏览文件 @ ffb72e42
......@@ -1113,13 +1113,13 @@ def tril(m, k=0):
Examples
--------
>>> import pytensor.tensor as pt
>>> pt.tril(pt.arange(1,13).reshape((4,3)), -1).eval()
>>> pt.tril(pt.arange(1, 13).reshape((4, 3)), -1).eval()
array([[ 0, 0, 0],
[ 4, 0, 0],
[ 7, 8, 0],
[10, 11, 12]])
>>> pt.tril(pt.arange(3*4*5).reshape((3, 4, 5))).eval()
>>> pt.tril(pt.arange(3 * 4 * 5).reshape((3, 4, 5))).eval()
array([[[ 0, 0, 0, 0, 0],
[ 5, 6, 0, 0, 0],
[10, 11, 12, 0, 0],
......@@ -1162,7 +1162,7 @@ def triu(m, k=0):
[ 0, 8, 9],
[ 0, 0, 12]])
>>> pt.triu(np.arange(3*4*5).reshape((3, 4, 5))).eval()
>>> pt.triu(np.arange(3 * 4 * 5).reshape((3, 4, 5))).eval()
array([[[ 0, 1, 2, 3, 4],
[ 0, 6, 7, 8, 9],
[ 0, 0, 12, 13, 14],
......@@ -2108,9 +2108,9 @@ class Split(COp):
>>> splits = pt.vector(dtype="int")
You have to declare right away how many split_points there will be.
>>> ra, rb, rc = pt.split(x, splits, n_splits = 3, axis = 0)
>>> ra, rb, rc = pt.split(x, splits, n_splits=3, axis=0)
>>> f = function([x, splits], [ra, rb, rc])
>>> a, b, c = f([0,1,2,3,4,5], [3, 2, 1])
>>> a, b, c = f([0, 1, 2, 3, 4, 5], [3, 2, 1])
>>> a
array([0, 1, 2])
>>> b
......@@ -3892,7 +3892,7 @@ def stacklists(arg):
>>> from pytensor.tensor import stacklists
>>> from pytensor.tensor.type import scalars, matrices
>>> from pytensor import function
>>> a, b, c, d = scalars('abcd')
>>> a, b, c, d = scalars("abcd")
>>> X = stacklists([[a, b], [c, d]])
>>> f = function([a, b, c, d], X)
>>> f(1, 2, 3, 4)
......@@ -3903,10 +3903,10 @@ def stacklists(arg):
a 2 by 2 grid:
>>> from numpy import ones
>>> a, b, c, d = matrices('abcd')
>>> a, b, c, d = matrices("abcd")
>>> X = stacklists([[a, b], [c, d]])
>>> f = function([a, b, c, d], X)
>>> x = ones((4, 4), 'float32')
>>> x = ones((4, 4), "float32")
>>> f(x, x, x, x).shape
(2, 2, 4, 4)
......
pytensor/tensor/einsum.py
浏览文件 @ ffb72e42
......@@ -467,6 +467,7 @@ def einsum(subscripts: str, *operands: "TensorLike", optimize=None) -> TensorVar
.. code-block:: python
import pytensor as pt
A = pt.matrix("A")
B = pt.matrix("B")
C = pt.einsum("ij, jk -> ik", A, B)
......@@ -481,6 +482,7 @@ def einsum(subscripts: str, *operands: "TensorLike", optimize=None) -> TensorVar
.. code-block:: python
import pytensor as pt
A = pt.tensor("A", shape=(None, 4, 5))
B = pt.tensor("B", shape=(None, 5, 6))
C = pt.einsum("bij, bjk -> bik", A, B)
......@@ -496,6 +498,7 @@ def einsum(subscripts: str, *operands: "TensorLike", optimize=None) -> TensorVar
.. code-block:: python
import pytensor as pt
A = pt.tensor("A", shape=(4, None, None, None, 5))
B = pt.tensor("B", shape=(5, None, None, None, 6))
C = pt.einsum("i...j, j...k -> ...ik", A, B)
......@@ -510,6 +513,7 @@ def einsum(subscripts: str, *operands: "TensorLike", optimize=None) -> TensorVar
.. code-block:: python
import pytensor as pt
x = pt.tensor("x", shape=(3,))
y = pt.tensor("y", shape=(4,))
z = pt.einsum("i, j -> ij", x, y)
......
pytensor/tensor/elemwise.py
浏览文件 @ ffb72e42
......@@ -77,7 +77,7 @@ class DimShuffle(ExternalCOp):
.. code-block:: python
DimShuffle((False, False, False), ['x', 2, 'x', 0, 1])
DimShuffle((False, False, False), ["x", 2, "x", 0, 1])
This `Op` will only work on 3d tensors with no broadcastable
dimensions. The first dimension will be broadcastable,
......@@ -101,16 +101,16 @@ class DimShuffle(ExternalCOp):
--------
.. code-block:: python
DimShuffle((), ['x']) # make a 0d (scalar) into a 1d vector
DimShuffle((), ["x"]) # make a 0d (scalar) into a 1d vector
DimShuffle((False, False), [0, 1]) # identity
DimShuffle((False, False), [1, 0]) # inverts the 1st and 2nd dimensions
DimShuffle((False,), ['x', 0]) # make a row out of a 1d vector
DimShuffle((False,), ["x", 0]) # make a row out of a 1d vector
# (N to 1xN)
DimShuffle((False,), [0, 'x']) # make a column out of a 1d vector
DimShuffle((False,), [0, "x"]) # make a column out of a 1d vector
# (N to Nx1)
DimShuffle((False, False, False), [2, 0, 1]) # AxBxC to CxAxB
DimShuffle((False, False), [0, 'x', 1]) # AxB to Ax1xB
DimShuffle((False, False), [1, 'x', 0]) # AxB to Bx1xA
DimShuffle((False, False), [0, "x", 1]) # AxB to Ax1xB
DimShuffle((False, False), [1, "x", 0]) # AxB to Bx1xA
The reordering of the dimensions can be done with the numpy.transpose
function.
......
pytensor/tensor/extra_ops.py
浏览文件 @ ffb72e42
......@@ -267,13 +267,13 @@ def searchsorted(x, v, side="left", sorter=None):
>>> from pytensor.tensor import extra_ops
>>> x = pt.dvector("x")
>>> idx = x.searchsorted(3)
>>> idx.eval({x: [1,2,3,4,5]})
>>> idx.eval({x: [1, 2, 3, 4, 5]})
array(2)
>>> extra_ops.searchsorted([1,2,3,4,5], 3).eval()
>>> extra_ops.searchsorted([1, 2, 3, 4, 5], 3).eval()
array(2)
>>> extra_ops.searchsorted([1,2,3,4,5], 3, side='right').eval()
>>> extra_ops.searchsorted([1, 2, 3, 4, 5], 3, side="right").eval()
array(3)
>>> extra_ops.searchsorted([1,2,3,4,5], [-10, 10, 2, 3]).eval()
>>> extra_ops.searchsorted([1, 2, 3, 4, 5], [-10, 10, 2, 3]).eval()
array([0, 5, 1, 2])
.. versionadded:: 0.9
......@@ -1176,7 +1176,7 @@ class Unique(Op):
>>> x = pytensor.tensor.vector()
>>> f = pytensor.function([x], Unique(True, True, False)(x))
>>> f([1, 2., 3, 4, 3, 2, 1.])
>>> f([1, 2.0, 3, 4, 3, 2, 1.0])
[array([1., 2., 3., 4.]), array([0, 1, 2, 3]), array([0, 1, 2, 3, 2, 1, 0])]
>>> y = pytensor.tensor.matrix()
......
pytensor/tensor/functional.py
浏览文件 @ ffb72e42
......@@ -39,9 +39,11 @@ def vectorize(func: Callable, signature: str | None = None) -> Callable:
import pytensor
import pytensor.tensor as pt
def func(x):
return pt.exp(x) / pt.sum(pt.exp(x))
vec_func = pt.vectorize(func, signature="(a)->(a)")
x = pt.matrix("x")
......@@ -58,9 +60,11 @@ def vectorize(func: Callable, signature: str | None = None) -> Callable:
import pytensor
import pytensor.tensor as pt
def func(x):
return x[0], x[-1]
vec_func = pt.vectorize(func, signature="(a)->(),()")
x = pt.matrix("x")
......
pytensor/tensor/io.py
浏览文件 @ ffb72e42
......@@ -80,8 +80,8 @@ def load(path, dtype, shape, mmap_mode=None):
--------
>>> from pytensor import *
>>> path = Variable(Generic(), None)
>>> x = tensor.load(path, 'int64', (None,))
>>> y = x*2
>>> x = tensor.load(path, "int64", (None,))
>>> y = x * 2
>>> fn = function([path], y)
>>> fn("stored-array.npy") # doctest: +SKIP
array([0, 2, 4, 6, 8], dtype=int64)
......
pytensor/tensor/math.py
浏览文件 @ ffb72e42
......@@ -2098,27 +2098,27 @@ def tensordot(
are compatible. The resulting tensor will have shape (2, 5, 6) -- the
dimensions that are not being summed:
>>> a = np.random.random((2,3,4))
>>> b = np.random.random((5,6,4,3))
>>> a = np.random.random((2, 3, 4))
>>> b = np.random.random((5, 6, 4, 3))
#tensordot
>>> c = np.tensordot(a, b, [[1,2],[3,2]])
>>> c = np.tensordot(a, b, [[1, 2], [3, 2]])
#loop replicating tensordot
>>> a0, a1, a2 = a.shape
>>> b0, b1, _, _ = b.shape
>>> cloop = np.zeros((a0,b0,b1))
>>> cloop = np.zeros((a0, b0, b1))
#loop over non-summed indices -- these exist
#in the tensor product.
>>> for i in range(a0):
... for j in range(b0):
... for k in range(b1):
... #loop over summed indices -- these don't exist
... #in the tensor product.
... # loop over summed indices -- these don't exist
... # in the tensor product.
... for l in range(a1):
... for m in range(a2):
... cloop[i,j,k] += a[i,l,m] * b[j,k,m,l]
... cloop[i, j, k] += a[i, l, m] * b[j, k, m, l]
>>> np.allclose(c, cloop)
True
......
pytensor/tensor/utils.py
浏览文件 @ ffb72e42
......@@ -61,8 +61,9 @@ def shape_of_variables(
--------
>>> import pytensor.tensor as pt
>>> from pytensor.graph.fg import FunctionGraph
>>> x = pt.matrix('x')
>>> y = x[512:]; y.name = 'y'
>>> x = pt.matrix("x")
>>> y = x[512:]
>>> y.name = "y"
>>> fgraph = FunctionGraph([x], [y], clone=False)
>>> d = shape_of_variables(fgraph, {x: (1024, 1024)})
>>> d[y]
......
Markdown 格式
0% 或
您添加了 0 人 到此讨论。请谨慎行事。
请先完成此评论的编辑！
请 注册 或者 后发表评论