galley-jl-python 0.1.0__tar.gz → 0.1.2__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {galley_jl_python-0.1.0 → galley_jl_python-0.1.2}/PKG-INFO +6 -6
- {galley_jl_python-0.1.0 → galley_jl_python-0.1.2}/README.md +5 -5
- {galley_jl_python-0.1.0 → galley_jl_python-0.1.2}/pyproject.toml +1 -1
- {galley_jl_python-0.1.0 → galley_jl_python-0.1.2}/src/galley_jl_python/__init__.py +6 -0
- {galley_jl_python-0.1.0 → galley_jl_python-0.1.2}/src/galley_jl_python/compiled.py +7 -7
- {galley_jl_python-0.1.0 → galley_jl_python-0.1.2}/src/galley_jl_python/julia.py +4 -5
- {galley_jl_python-0.1.0 → galley_jl_python-0.1.2}/src/galley_jl_python/juliapkg.json +6 -3
- {galley_jl_python-0.1.0 → galley_jl_python-0.1.2}/src/galley_jl_python/juliapkg_dev.json +3 -3
- {galley_jl_python-0.1.0 → galley_jl_python-0.1.2}/src/galley_jl_python/linalg/_linalg.py +1 -1
- {galley_jl_python-0.1.0 → galley_jl_python-0.1.2}/src/galley_jl_python/tensor.py +184 -45
- {galley_jl_python-0.1.0 → galley_jl_python-0.1.2}/LICENSE +0 -0
- {galley_jl_python-0.1.0 → galley_jl_python-0.1.2}/src/galley_jl_python/_array_api_info.py +0 -0
- {galley_jl_python-0.1.0 → galley_jl_python-0.1.2}/src/galley_jl_python/_sysimage.py +0 -0
- {galley_jl_python-0.1.0 → galley_jl_python-0.1.2}/src/galley_jl_python/dtypes.py +0 -0
- {galley_jl_python-0.1.0 → galley_jl_python-0.1.2}/src/galley_jl_python/einstein.py +0 -0
- {galley_jl_python-0.1.0 → galley_jl_python-0.1.2}/src/galley_jl_python/errors.py +0 -0
- {galley_jl_python-0.1.0 → galley_jl_python-0.1.2}/src/galley_jl_python/finch_fused.py +0 -0
- {galley_jl_python-0.1.0 → galley_jl_python-0.1.2}/src/galley_jl_python/io.py +0 -0
- {galley_jl_python-0.1.0 → galley_jl_python-0.1.2}/src/galley_jl_python/levels.py +0 -0
- {galley_jl_python-0.1.0 → galley_jl_python-0.1.2}/src/galley_jl_python/linalg/__init__.py +0 -0
- {galley_jl_python-0.1.0 → galley_jl_python-0.1.2}/src/galley_jl_python/typing.py +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: galley-jl-python
|
|
3
|
-
Version: 0.1.
|
|
3
|
+
Version: 0.1.2
|
|
4
4
|
Summary: Sparse Tensor Programming in Python powered by Finch.jl
|
|
5
5
|
License-File: LICENSE
|
|
6
6
|
Author: Willow Ahrens
|
|
@@ -79,8 +79,8 @@ Run any other command inside an environment with `pixi run`, for example
|
|
|
79
79
|
The package is still built and published with Poetry (see
|
|
80
80
|
[Publishing](#publishing)).
|
|
81
81
|
|
|
82
|
-
### Working with a local copy of
|
|
83
|
-
The `develop.py` script can be used to set up a local copy of
|
|
82
|
+
### Working with a local copy of Galley.jl
|
|
83
|
+
The `develop.py` script can be used to set up a local copy of Galley.jl for
|
|
84
84
|
development. Run it with `pixi run python develop.py`.
|
|
85
85
|
|
|
86
86
|
```
|
|
@@ -89,12 +89,12 @@ Usage:
|
|
|
89
89
|
|
|
90
90
|
Options:
|
|
91
91
|
--restore Restore the original juliapkg.json file.
|
|
92
|
-
--path Path to the local copy of
|
|
92
|
+
--path Path to the local copy of Galley.jl [default: ../Galley.jl].
|
|
93
93
|
```
|
|
94
94
|
|
|
95
95
|
### Julia sysimage
|
|
96
96
|
|
|
97
|
-
Most of Galley's startup time is Julia compiling
|
|
97
|
+
Most of Galley's startup time is Julia compiling Galley.jl itself. A prebuilt Julia
|
|
98
98
|
sysimage removes it: the first matmul of a session drops from about 4.5 minutes
|
|
99
99
|
to a few seconds.
|
|
100
100
|
|
|
@@ -112,7 +112,7 @@ to a few seconds.
|
|
|
112
112
|
An image works only with the exact Julia and package versions it was built
|
|
113
113
|
from. For this reason `src/galley_jl_python/juliapkg.json` pins Julia and every
|
|
114
114
|
Julia package. Each image's name includes a hash of that environment. A local
|
|
115
|
-
|
|
115
|
+
Galley.jl from `develop.py` therefore runs without the image.
|
|
116
116
|
|
|
117
117
|
To update the Julia dependencies:
|
|
118
118
|
|
|
@@ -48,8 +48,8 @@ Run any other command inside an environment with `pixi run`, for example
|
|
|
48
48
|
The package is still built and published with Poetry (see
|
|
49
49
|
[Publishing](#publishing)).
|
|
50
50
|
|
|
51
|
-
### Working with a local copy of
|
|
52
|
-
The `develop.py` script can be used to set up a local copy of
|
|
51
|
+
### Working with a local copy of Galley.jl
|
|
52
|
+
The `develop.py` script can be used to set up a local copy of Galley.jl for
|
|
53
53
|
development. Run it with `pixi run python develop.py`.
|
|
54
54
|
|
|
55
55
|
```
|
|
@@ -58,12 +58,12 @@ Usage:
|
|
|
58
58
|
|
|
59
59
|
Options:
|
|
60
60
|
--restore Restore the original juliapkg.json file.
|
|
61
|
-
--path Path to the local copy of
|
|
61
|
+
--path Path to the local copy of Galley.jl [default: ../Galley.jl].
|
|
62
62
|
```
|
|
63
63
|
|
|
64
64
|
### Julia sysimage
|
|
65
65
|
|
|
66
|
-
Most of Galley's startup time is Julia compiling
|
|
66
|
+
Most of Galley's startup time is Julia compiling Galley.jl itself. A prebuilt Julia
|
|
67
67
|
sysimage removes it: the first matmul of a session drops from about 4.5 minutes
|
|
68
68
|
to a few seconds.
|
|
69
69
|
|
|
@@ -81,7 +81,7 @@ to a few seconds.
|
|
|
81
81
|
An image works only with the exact Julia and package versions it was built
|
|
82
82
|
from. For this reason `src/galley_jl_python/juliapkg.json` pins Julia and every
|
|
83
83
|
Julia package. Each image's name includes a hash of that environment. A local
|
|
84
|
-
|
|
84
|
+
Galley.jl from `develop.py` therefore runs without the image.
|
|
85
85
|
|
|
86
86
|
To update the Julia dependencies:
|
|
87
87
|
|
|
@@ -177,11 +177,14 @@ from .tensor import (
|
|
|
177
177
|
log10,
|
|
178
178
|
logaddexp,
|
|
179
179
|
logical_and,
|
|
180
|
+
logical_not,
|
|
180
181
|
logical_or,
|
|
181
182
|
logical_xor,
|
|
182
183
|
max,
|
|
184
|
+
maximum,
|
|
183
185
|
mean,
|
|
184
186
|
min,
|
|
187
|
+
minimum,
|
|
185
188
|
moveaxis,
|
|
186
189
|
nonzero,
|
|
187
190
|
ones,
|
|
@@ -305,12 +308,15 @@ __all__ = [
|
|
|
305
308
|
"log10",
|
|
306
309
|
"logaddexp",
|
|
307
310
|
"logical_and",
|
|
311
|
+
"logical_not",
|
|
308
312
|
"logical_or",
|
|
309
313
|
"logical_xor",
|
|
310
314
|
"matmul",
|
|
311
315
|
"max",
|
|
316
|
+
"maximum",
|
|
312
317
|
"mean",
|
|
313
318
|
"min",
|
|
319
|
+
"minimum",
|
|
314
320
|
"moveaxis",
|
|
315
321
|
"multiply",
|
|
316
322
|
"nan",
|
|
@@ -73,7 +73,7 @@ def _recurse_iter_compute(x: IterObj, /, *, compute_kwargs: dict[str, Any]) -> I
|
|
|
73
73
|
jl_len = len(jl_computed)
|
|
74
74
|
# This doesn't return an iterable of arrays -- only a single array
|
|
75
75
|
# for `len(jl_computed) == 1`
|
|
76
|
-
jl_computed = jl.
|
|
76
|
+
jl_computed = jl.Galley.compute(*jl_computed, **compute_kwargs)
|
|
77
77
|
if jl_len == 1:
|
|
78
78
|
jl_computed = (jl_computed,)
|
|
79
79
|
|
|
@@ -129,16 +129,16 @@ class AbstractScheduler:
|
|
|
129
129
|
|
|
130
130
|
class GalleyScheduler(AbstractScheduler):
|
|
131
131
|
def get_julia_scheduler(self) -> JuliaObj:
|
|
132
|
-
return jl.
|
|
132
|
+
return jl.Galley.galley_scheduler(verbose=self.verbose)
|
|
133
133
|
|
|
134
134
|
|
|
135
135
|
class DefaultScheduler(AbstractScheduler):
|
|
136
136
|
def get_julia_scheduler(self) -> JuliaObj:
|
|
137
|
-
return jl.
|
|
137
|
+
return jl.Galley.default_scheduler(verbose=self.verbose)
|
|
138
138
|
|
|
139
139
|
|
|
140
140
|
def set_optimizer(opt: AbstractScheduler) -> None:
|
|
141
|
-
jl.
|
|
141
|
+
jl.Galley.set_scheduler_b(opt.get_julia_scheduler())
|
|
142
142
|
|
|
143
143
|
|
|
144
144
|
set_optimizer(GalleyScheduler())
|
|
@@ -148,7 +148,7 @@ def lazy(tensor: Tensor) -> Tensor:
|
|
|
148
148
|
from .tensor import Tensor
|
|
149
149
|
|
|
150
150
|
if tensor.is_computed():
|
|
151
|
-
return Tensor(jl.
|
|
151
|
+
return Tensor(jl.Galley.LazyTensor(tensor._obj))
|
|
152
152
|
return tensor
|
|
153
153
|
|
|
154
154
|
|
|
@@ -159,9 +159,9 @@ def compute(
|
|
|
159
159
|
|
|
160
160
|
if not tensor.is_computed():
|
|
161
161
|
if opt is None:
|
|
162
|
-
return Tensor(jl.
|
|
162
|
+
return Tensor(jl.Galley.compute(tensor._obj, tag=tag))
|
|
163
163
|
return Tensor(
|
|
164
|
-
jl.
|
|
164
|
+
jl.Galley.compute(
|
|
165
165
|
tensor._obj,
|
|
166
166
|
verbose=opt.verbose,
|
|
167
167
|
ctx=opt.get_julia_scheduler(),
|
|
@@ -19,17 +19,16 @@ from . import _sysimage # noqa: E402
|
|
|
19
19
|
|
|
20
20
|
_sysimage.configure()
|
|
21
21
|
|
|
22
|
-
# To change the version of
|
|
22
|
+
# To change the version of Galley used, see the documentation for pyjuliapkg here: https://github.com/JuliaPy/pyjuliapkg
|
|
23
23
|
# Use pyjuliapkg to modify the `juliapkg.json` file in the root of this repo.
|
|
24
|
-
# You can also run `develop.py` to quickly use a local copy of
|
|
24
|
+
# You can also run `develop.py` to quickly use a local copy of Galley.jl.
|
|
25
25
|
# An example development json is found in `juliapkg_dev.json`
|
|
26
26
|
# Julia and all packages are pinned so the prebuilt sysimage matches; update the
|
|
27
|
-
# pins with `scripts/sysimage/pin_julia_deps.py`.
|
|
28
|
-
# commit (`rev`), which that script leaves alone.
|
|
27
|
+
# pins with `scripts/sysimage/pin_julia_deps.py`.
|
|
29
28
|
import juliacall as jc # noqa: E402, F401
|
|
30
29
|
from juliacall import Main as jl # noqa: E402, F401
|
|
31
30
|
|
|
32
|
-
jl.seval("using
|
|
31
|
+
jl.seval("using Galley")
|
|
33
32
|
jl.seval("using HDF5")
|
|
34
33
|
jl.seval("using NPZ")
|
|
35
34
|
jl.seval("using TensorMarket")
|
|
@@ -82,14 +82,17 @@
|
|
|
82
82
|
"version": "=1.17.1"
|
|
83
83
|
},
|
|
84
84
|
"Finch": {
|
|
85
|
-
"
|
|
86
|
-
"
|
|
87
|
-
"uuid": "9177782c-1635-4eb9-9bfb-d9dfa25e6bce"
|
|
85
|
+
"uuid": "9177782c-1635-4eb9-9bfb-d9dfa25e6bce",
|
|
86
|
+
"version": "=1.8.0"
|
|
88
87
|
},
|
|
89
88
|
"GPUArraysCore": {
|
|
90
89
|
"uuid": "46192b85-c4d5-4398-a991-12ede77f4527",
|
|
91
90
|
"version": "=0.2.1"
|
|
92
91
|
},
|
|
92
|
+
"Galley": {
|
|
93
|
+
"uuid": "0b99cc74-af0a-4786-be99-594ac14242a4",
|
|
94
|
+
"version": "=0.2.2"
|
|
95
|
+
},
|
|
93
96
|
"Gamma": {
|
|
94
97
|
"uuid": "a0844989-3bd2-4988-8bea-c9407ab0941b",
|
|
95
98
|
"version": "=1.2.0"
|
|
@@ -1,10 +1,10 @@
|
|
|
1
1
|
{
|
|
2
2
|
"julia": "1.10 - 1.12",
|
|
3
3
|
"packages": {
|
|
4
|
-
"
|
|
4
|
+
"Galley": {
|
|
5
5
|
"dev": true,
|
|
6
|
-
"path": "../
|
|
7
|
-
"uuid": "
|
|
6
|
+
"path": "../Galley.jl",
|
|
7
|
+
"uuid": "0b99cc74-af0a-4786-be99-594ac14242a4"
|
|
8
8
|
},
|
|
9
9
|
"HDF5": {
|
|
10
10
|
"uuid": "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f",
|
|
@@ -25,9 +25,9 @@ from .levels import (
|
|
|
25
25
|
)
|
|
26
26
|
from .typing import Device, DType, JuliaObj, OrderType, TupleOf3Arrays, spmatrix
|
|
27
27
|
|
|
28
|
-
# Builds a 0-d
|
|
28
|
+
# Builds a 0-d Galley tensor in Julia so the element keeps its Julia eltype.
|
|
29
29
|
_scalar_to_0d = jl.seval(
|
|
30
|
-
"(T, fv, x) ->
|
|
30
|
+
"(T, fv, x) -> Galley.Tensor(Galley.Element{convert(T, fv)}(T[convert(T, x)]))"
|
|
31
31
|
)
|
|
32
32
|
|
|
33
33
|
|
|
@@ -37,9 +37,65 @@ class SparseArray:
|
|
|
37
37
|
"""
|
|
38
38
|
|
|
39
39
|
|
|
40
|
+
_WEAK_SCALAR_TYPE = """
|
|
41
|
+
galley_weak_scalar_type(::Type{Bool}, ::Type{Bool}) = Bool
|
|
42
|
+
galley_weak_scalar_type(::Type{T}, ::Type{Bool}) where {T} = T
|
|
43
|
+
galley_weak_scalar_type(::Type{Bool}, ::Type{S}) where {S<:Integer} = S
|
|
44
|
+
galley_weak_scalar_type(::Type{T}, ::Type{<:Integer}) where {T<:Number} = T
|
|
45
|
+
galley_weak_scalar_type(
|
|
46
|
+
::Type{T}, ::Type{<:AbstractFloat}
|
|
47
|
+
) where {T<:Union{AbstractFloat,Complex}} = T
|
|
48
|
+
galley_weak_scalar_type(::Type{T}, ::Type{<:Complex}) where {T<:Complex} = T
|
|
49
|
+
galley_weak_scalar_type(
|
|
50
|
+
::Type{T}, ::Type{<:Complex}
|
|
51
|
+
) where {T<:AbstractFloat} = Complex{T}
|
|
52
|
+
galley_weak_scalar_type(::Type, ::Type{S}) where {S} = S
|
|
53
|
+
"""
|
|
54
|
+
|
|
55
|
+
|
|
56
|
+
def _weak_scalar(element_type, value):
|
|
57
|
+
"""A Python scalar operand as the Array API's weak scalars are: of the
|
|
58
|
+
array's type when it is of the same kind (an int with an integer array, a
|
|
59
|
+
float with a floating one), and otherwise of the default type (a float
|
|
60
|
+
with a Boolean array is a Float64). Reads lazy tensors' types too."""
|
|
61
|
+
if not hasattr(jl, "galley_weak_scalar_type"):
|
|
62
|
+
jl.seval(_WEAK_SCALAR_TYPE)
|
|
63
|
+
if isinstance(value, np.generic):
|
|
64
|
+
value = value.item() # Julia reads np.bool_ as a Python object.
|
|
65
|
+
value_type = jl.typeof(jc.convert(jl.Any, value))
|
|
66
|
+
# jc.convert keeps the result in Julia; a Julia number returned to Python
|
|
67
|
+
# becomes a Python number, and so a Float64, again.
|
|
68
|
+
return jc.convert(jl.galley_weak_scalar_type(element_type, value_type), value)
|
|
69
|
+
|
|
70
|
+
|
|
71
|
+
def _scalar_tensor(value) -> Tensor:
|
|
72
|
+
"""A Julia scalar as a 0-d tensor whose fill is the scalar itself."""
|
|
73
|
+
return Tensor(jl.swizzle(_scalar_to_0d(jl.typeof(value), value, value)))
|
|
74
|
+
|
|
75
|
+
|
|
76
|
+
# `where` picks with `ifelse` after promoting both branches to one type, as
|
|
77
|
+
# the Array API requires. Julia's `ifelse` returns whichever branch it picks
|
|
78
|
+
# unchanged, so on its own it would give a Union{Float64, Int64} result for a
|
|
79
|
+
# float branch and an integer one.
|
|
80
|
+
_where_op = jl.seval("(c, a, b) -> ifelse(c, promote(a, b)...)")
|
|
81
|
+
|
|
82
|
+
|
|
83
|
+
def _element_fill(dtype, fill_value):
|
|
84
|
+
"""fill_value as an Element level's fill for values of dtype.
|
|
85
|
+
|
|
86
|
+
Finch requires an isbits fill. Julia converts NumPy's numeric scalars to
|
|
87
|
+
Julia numbers, but not `np.bool_` or 0-d arrays (which `sparse.COO` uses for
|
|
88
|
+
its fill), so both are normalized first, and Booleans become Python bools.
|
|
89
|
+
"""
|
|
90
|
+
fill_value = np.asarray(fill_value).item()
|
|
91
|
+
if np.dtype(dtype) == np.bool_:
|
|
92
|
+
return bool(fill_value)
|
|
93
|
+
return np.dtype(dtype).type(fill_value)
|
|
94
|
+
|
|
95
|
+
|
|
40
96
|
class Tensor(_Display, SparseArray):
|
|
41
97
|
"""
|
|
42
|
-
A wrapper class for
|
|
98
|
+
A wrapper class for Galley.Tensor and Galley.SwizzleArray.
|
|
43
99
|
|
|
44
100
|
Constructors
|
|
45
101
|
------------
|
|
@@ -58,7 +114,7 @@ class Tensor(_Display, SparseArray):
|
|
|
58
114
|
|
|
59
115
|
Parameters
|
|
60
116
|
----------
|
|
61
|
-
obj : np.ndarray or scipy.sparse or Storage or
|
|
117
|
+
obj : np.ndarray or scipy.sparse or Storage or Galley.SwizzleArray
|
|
62
118
|
Input to construct a Tensor. It's a no-copy operation of for NumPy and
|
|
63
119
|
SciPy input. For Storage it's levels' description with order. The order
|
|
64
120
|
numbers the dimensions from the fastest to slowest. The leaf nodes have
|
|
@@ -75,7 +131,7 @@ class Tensor(_Display, SparseArray):
|
|
|
75
131
|
Returns
|
|
76
132
|
-------
|
|
77
133
|
Tensor
|
|
78
|
-
Python wrapper for
|
|
134
|
+
Python wrapper for a Galley Tensor.
|
|
79
135
|
|
|
80
136
|
Examples
|
|
81
137
|
--------
|
|
@@ -108,10 +164,10 @@ class Tensor(_Display, SparseArray):
|
|
|
108
164
|
fill_value: np.number | None = None,
|
|
109
165
|
copy: bool | None = None,
|
|
110
166
|
):
|
|
111
|
-
if isinstance(obj, int | float | complex | bool | list):
|
|
167
|
+
if isinstance(obj, int | float | complex | bool | list | tuple):
|
|
112
168
|
if copy is False:
|
|
113
169
|
raise ValueError(
|
|
114
|
-
"copy=False isn't supported for scalar inputs and Python
|
|
170
|
+
"copy=False isn't supported for scalar inputs and Python sequences"
|
|
115
171
|
)
|
|
116
172
|
obj = np.asarray(obj)
|
|
117
173
|
if fill_value is None:
|
|
@@ -130,11 +186,11 @@ class Tensor(_Display, SparseArray):
|
|
|
130
186
|
obj.order, self.get_lvl_ndim(obj.levels_descr._obj)
|
|
131
187
|
)
|
|
132
188
|
self._obj = jl.swizzle(jl.Tensor(obj.levels_descr._obj), *order)
|
|
133
|
-
elif jl.isa(obj, jl.
|
|
189
|
+
elif jl.isa(obj, jl.Galley.Tensor): # raw-Julia-object constructors
|
|
134
190
|
if copy:
|
|
135
191
|
self._raise_julia_copy_not_supported()
|
|
136
192
|
self._obj = jl.swizzle(obj, *tuple(range(1, jl.ndims(obj) + 1)))
|
|
137
|
-
elif jl.isa(obj, jl.
|
|
193
|
+
elif jl.isa(obj, jl.Galley.SwizzleArray) or jl.isa(obj, jl.Galley.LazyTensor):
|
|
138
194
|
if copy:
|
|
139
195
|
self._raise_julia_copy_not_supported()
|
|
140
196
|
self._obj = obj
|
|
@@ -165,14 +221,51 @@ class Tensor(_Display, SparseArray):
|
|
|
165
221
|
return self._elemwise_op("/", other)
|
|
166
222
|
|
|
167
223
|
def __floordiv__(self, other):
|
|
168
|
-
return self._elemwise_op("
|
|
224
|
+
return self._elemwise_op("Galley.fld_nothrow", other)
|
|
169
225
|
|
|
170
226
|
def __mod__(self, other):
|
|
171
|
-
return self._elemwise_op("
|
|
227
|
+
return self._elemwise_op("Galley.mod_nothrow", other)
|
|
172
228
|
|
|
173
229
|
def __pow__(self, other):
|
|
174
230
|
return self._elemwise_op("^", other)
|
|
175
231
|
|
|
232
|
+
# Reflected operators, for a scalar on the left (`1 - x`, `2 ** x`).
|
|
233
|
+
def __radd__(self, other):
|
|
234
|
+
return self._elemwise_op("+", other, reflected=True)
|
|
235
|
+
|
|
236
|
+
def __rsub__(self, other):
|
|
237
|
+
return self._elemwise_op("-", other, reflected=True)
|
|
238
|
+
|
|
239
|
+
def __rmul__(self, other):
|
|
240
|
+
return self._elemwise_op("*", other, reflected=True)
|
|
241
|
+
|
|
242
|
+
def __rtruediv__(self, other):
|
|
243
|
+
return self._elemwise_op("/", other, reflected=True)
|
|
244
|
+
|
|
245
|
+
def __rfloordiv__(self, other):
|
|
246
|
+
return self._elemwise_op("Galley.fld_nothrow", other, reflected=True)
|
|
247
|
+
|
|
248
|
+
def __rmod__(self, other):
|
|
249
|
+
return self._elemwise_op("Galley.mod_nothrow", other, reflected=True)
|
|
250
|
+
|
|
251
|
+
def __rpow__(self, other):
|
|
252
|
+
return self._elemwise_op("^", other, reflected=True)
|
|
253
|
+
|
|
254
|
+
def __rand__(self, other):
|
|
255
|
+
return self._elemwise_op("&", other, reflected=True)
|
|
256
|
+
|
|
257
|
+
def __ror__(self, other):
|
|
258
|
+
return self._elemwise_op("|", other, reflected=True)
|
|
259
|
+
|
|
260
|
+
def __rxor__(self, other):
|
|
261
|
+
return self._elemwise_op("xor", other, reflected=True)
|
|
262
|
+
|
|
263
|
+
def __rlshift__(self, other):
|
|
264
|
+
return self._elemwise_op("<<", other, reflected=True)
|
|
265
|
+
|
|
266
|
+
def __rrshift__(self, other):
|
|
267
|
+
return self._elemwise_op(">>", other, reflected=True)
|
|
268
|
+
|
|
176
269
|
@compiled()
|
|
177
270
|
def __matmul__(self, other: Tensor) -> Tensor:
|
|
178
271
|
if self.ndim == 0 or other.ndim == 0:
|
|
@@ -227,20 +320,21 @@ class Tensor(_Display, SparseArray):
|
|
|
227
320
|
def __ne__(self, other):
|
|
228
321
|
return self._elemwise_op("!=", other)
|
|
229
322
|
|
|
230
|
-
def _elemwise_op(
|
|
323
|
+
def _elemwise_op(
|
|
324
|
+
self, op: str, other: Tensor | None = None, *, reflected: bool = False
|
|
325
|
+
) -> Tensor:
|
|
326
|
+
fn = jl.seval(op)
|
|
231
327
|
if other is None:
|
|
232
|
-
result = jl.broadcast(
|
|
328
|
+
result = jl.broadcast(fn, self._obj)
|
|
233
329
|
else:
|
|
234
330
|
if np.isscalar(other):
|
|
235
|
-
other =
|
|
331
|
+
other = _weak_scalar(jl.eltype(self._obj), other)
|
|
236
332
|
else:
|
|
237
333
|
other = jl.permutedims(other._obj, tuple(range(other.ndim, 0, -1)))
|
|
238
334
|
# inverse swizzle, so `broadcast` appends new dims to the front
|
|
239
|
-
|
|
240
|
-
|
|
241
|
-
|
|
242
|
-
other,
|
|
243
|
-
)
|
|
335
|
+
this = jl.permutedims(self._obj, tuple(range(self.ndim, 0, -1)))
|
|
336
|
+
args = (other, this) if reflected else (this, other)
|
|
337
|
+
result = jl.broadcast(fn, *args)
|
|
244
338
|
# swizzle back to the original order
|
|
245
339
|
result = jl.permutedims(result, tuple(range(jl.ndims(result), 0, -1)))
|
|
246
340
|
|
|
@@ -280,9 +374,11 @@ class Tensor(_Display, SparseArray):
|
|
|
280
374
|
key = _add_plus_one(key, self.shape)
|
|
281
375
|
|
|
282
376
|
result = self._obj[key]
|
|
283
|
-
if jl.isa(result, jl.
|
|
377
|
+
if jl.isa(result, jl.Galley.SwizzleArray) or jl.isa(
|
|
378
|
+
result, jl.Galley.LazyTensor
|
|
379
|
+
):
|
|
284
380
|
return Tensor(result)
|
|
285
|
-
if jl.isa(result, jl.
|
|
381
|
+
if jl.isa(result, jl.Galley.Tensor):
|
|
286
382
|
return Tensor(jl.swizzle(result, *range(1, jl.ndims(result) + 1)))
|
|
287
383
|
# the Array API requires a 0-d array, not a scalar, for integer-only keys
|
|
288
384
|
return Tensor(jl.swizzle(_scalar_to_0d(self.dtype, self.fill_value, result)))
|
|
@@ -311,7 +407,7 @@ class Tensor(_Display, SparseArray):
|
|
|
311
407
|
def _is_dense(self) -> bool:
|
|
312
408
|
lvl = self._obj.body.lvl
|
|
313
409
|
for _ in self.shape:
|
|
314
|
-
if not jl.isa(lvl, jl.
|
|
410
|
+
if not jl.isa(lvl, jl.Galley.Dense):
|
|
315
411
|
return False
|
|
316
412
|
lvl = lvl.lvl
|
|
317
413
|
return True
|
|
@@ -320,6 +416,12 @@ class Tensor(_Display, SparseArray):
|
|
|
320
416
|
def _order(self) -> tuple[int, ...]:
|
|
321
417
|
return jl.typeof(self._obj).parameters[1]
|
|
322
418
|
|
|
419
|
+
@property
|
|
420
|
+
def T(self) -> Tensor:
|
|
421
|
+
if self.ndim != 2:
|
|
422
|
+
raise ValueError(f"`T` requires a 2-dimensional array, got {self.ndim=}.")
|
|
423
|
+
return self.permute_dims((1, 0))
|
|
424
|
+
|
|
323
425
|
@property
|
|
324
426
|
def mT(self) -> Tensor:
|
|
325
427
|
axes = list(range(self.ndim))
|
|
@@ -340,7 +442,7 @@ class Tensor(_Display, SparseArray):
|
|
|
340
442
|
return self
|
|
341
443
|
|
|
342
444
|
def is_computed(self) -> bool:
|
|
343
|
-
return not jl.isa(self._obj, jl.
|
|
445
|
+
return not jl.isa(self._obj, jl.Galley.LazyTensor)
|
|
344
446
|
|
|
345
447
|
@classmethod
|
|
346
448
|
def preprocess_order(cls, order: OrderType, ndim: int) -> tuple[int, ...]:
|
|
@@ -370,7 +472,7 @@ class Tensor(_Display, SparseArray):
|
|
|
370
472
|
while True:
|
|
371
473
|
ndim += 1
|
|
372
474
|
lvl = lvl.lvl
|
|
373
|
-
if jl.isa(lvl, jl.
|
|
475
|
+
if jl.isa(lvl, jl.Galley.Element):
|
|
374
476
|
break
|
|
375
477
|
return ndim
|
|
376
478
|
|
|
@@ -433,12 +535,7 @@ class Tensor(_Display, SparseArray):
|
|
|
433
535
|
order = cls.preprocess_order(order_char, arr.ndim)
|
|
434
536
|
inv_order = tuple(i - 1 for i in jl.invperm(order))
|
|
435
537
|
|
|
436
|
-
|
|
437
|
-
if (
|
|
438
|
-
dtype == np.bool_
|
|
439
|
-
): # Fails with: Finch currently only supports isbits defaults
|
|
440
|
-
dtype = jl_dtypes.bool
|
|
441
|
-
fill_value = dtype(fill_value)
|
|
538
|
+
fill_value = _element_fill(arr.dtype, fill_value)
|
|
442
539
|
lvl = Element(fill_value, arr.reshape(-1, order=order_char))
|
|
443
540
|
for i in inv_order:
|
|
444
541
|
lvl = Dense(lvl, arr.shape[i])
|
|
@@ -503,11 +600,10 @@ class Tensor(_Display, SparseArray):
|
|
|
503
600
|
def construct_coo_jl_object(
|
|
504
601
|
cls, coords, data, shape, order, fill_value=0.0
|
|
505
602
|
) -> JuliaObj:
|
|
506
|
-
assert len(coords) == 2
|
|
507
603
|
ndim = len(shape)
|
|
508
604
|
order = cls.preprocess_order(order, ndim)
|
|
509
605
|
|
|
510
|
-
lvl = jl.Element(data.dtype
|
|
606
|
+
lvl = jl.Element(_element_fill(data.dtype, fill_value), data)
|
|
511
607
|
ptr = jl.Vector[jl.Int]([1, len(data) + 1])
|
|
512
608
|
tbl = tuple(jl.PlusOneVector(arr) for arr in coords)
|
|
513
609
|
|
|
@@ -532,11 +628,10 @@ class Tensor(_Display, SparseArray):
|
|
|
532
628
|
assert len(shape) == 2
|
|
533
629
|
|
|
534
630
|
data, indices, indptr = arg
|
|
535
|
-
dtype = data.dtype.type
|
|
536
631
|
indices = jl.PlusOneVector(indices)
|
|
537
632
|
indptr = jl.PlusOneVector(indptr)
|
|
538
633
|
|
|
539
|
-
lvl = jl.Element(dtype
|
|
634
|
+
lvl = jl.Element(_element_fill(data.dtype, fill_value), data)
|
|
540
635
|
return jl.swizzle(
|
|
541
636
|
jl.Tensor(
|
|
542
637
|
jl.Dense(jl.SparseList(lvl, shape[0], indptr, indices), shape[1])
|
|
@@ -579,7 +674,6 @@ class Tensor(_Display, SparseArray):
|
|
|
579
674
|
assert isinstance(arg, tuple) and len(arg) == 3
|
|
580
675
|
|
|
581
676
|
data, indices_list, indptr_list = arg
|
|
582
|
-
dtype = data.dtype.type
|
|
583
677
|
|
|
584
678
|
assert len(indices_list) == len(shape) - 1
|
|
585
679
|
assert len(indptr_list) == len(shape) - 1
|
|
@@ -587,7 +681,7 @@ class Tensor(_Display, SparseArray):
|
|
|
587
681
|
indices_list = [jl.PlusOneVector(i) for i in indices_list]
|
|
588
682
|
indptr_list = [jl.PlusOneVector(i) for i in indptr_list]
|
|
589
683
|
|
|
590
|
-
lvl = jl.Element(dtype
|
|
684
|
+
lvl = jl.Element(_element_fill(data.dtype, fill_value), data)
|
|
591
685
|
for size, indices, indptr in zip(
|
|
592
686
|
shape[:-1], indices_list, indptr_list, strict=False
|
|
593
687
|
):
|
|
@@ -906,7 +1000,20 @@ def astype(x: Tensor, dtype: DType, /, *, copy: bool = True) -> Tensor:
|
|
|
906
1000
|
return Tensor(jl.swizzle(result, *x.get_order(zero_indexing=False)))
|
|
907
1001
|
|
|
908
1002
|
|
|
909
|
-
def where(
|
|
1003
|
+
def where(
|
|
1004
|
+
condition: Tensor,
|
|
1005
|
+
x1: Tensor | complex,
|
|
1006
|
+
x2: Tensor | complex,
|
|
1007
|
+
/,
|
|
1008
|
+
) -> Tensor:
|
|
1009
|
+
# A Python scalar takes the other branch's type when it is of the same
|
|
1010
|
+
# kind, as the Array API's weak scalars do.
|
|
1011
|
+
if np.isscalar(x1) and np.isscalar(x2):
|
|
1012
|
+
raise TypeError("where needs at least one of x1 and x2 to be an array.")
|
|
1013
|
+
if np.isscalar(x1):
|
|
1014
|
+
x1 = _scalar_tensor(_weak_scalar(jl.eltype(x2._obj), x1))
|
|
1015
|
+
if np.isscalar(x2):
|
|
1016
|
+
x2 = _scalar_tensor(_weak_scalar(jl.eltype(x1._obj), x2))
|
|
910
1017
|
axis_cond, axis_x1, axis_x2 = (
|
|
911
1018
|
range(condition.ndim, 0, -1),
|
|
912
1019
|
range(x1.ndim, 0, -1),
|
|
@@ -914,7 +1021,7 @@ def where(condition: Tensor, x1: Tensor, x2: Tensor, /) -> Tensor:
|
|
|
914
1021
|
)
|
|
915
1022
|
# inverse swizzle, so `broadcast` appends new dims to the front
|
|
916
1023
|
result = jl.broadcast(
|
|
917
|
-
|
|
1024
|
+
_where_op,
|
|
918
1025
|
jl.permutedims(condition._obj, tuple(axis_cond)),
|
|
919
1026
|
jl.permutedims(x1._obj, tuple(axis_x1)),
|
|
920
1027
|
jl.permutedims(x2._obj, tuple(axis_x2)),
|
|
@@ -965,7 +1072,7 @@ def _reduce_sum_prod(
|
|
|
965
1072
|
|
|
966
1073
|
result = Tensor(result)
|
|
967
1074
|
|
|
968
|
-
if jl.isa(result._obj, jl.
|
|
1075
|
+
if jl.isa(result._obj, jl.Galley.LazyTensor):
|
|
969
1076
|
if dtype is not None:
|
|
970
1077
|
raise ValueError(
|
|
971
1078
|
"`dtype` keyword for `sum` and `prod` in the lazy mode isn't supported"
|
|
@@ -1102,7 +1209,7 @@ def argmin(
|
|
|
1102
1209
|
axis: int | None = None,
|
|
1103
1210
|
keepdims: bool = False,
|
|
1104
1211
|
) -> Tensor:
|
|
1105
|
-
return _reduce(x, jl.
|
|
1212
|
+
return _reduce(x, jl.Galley.argmin_python, axis, keepdims)
|
|
1106
1213
|
|
|
1107
1214
|
|
|
1108
1215
|
def argmax(
|
|
@@ -1112,7 +1219,7 @@ def argmax(
|
|
|
1112
1219
|
axis: int | None = None,
|
|
1113
1220
|
keepdims: bool = False,
|
|
1114
1221
|
) -> Tensor:
|
|
1115
|
-
return _reduce(x, jl.
|
|
1222
|
+
return _reduce(x, jl.Galley.argmax_python, axis, keepdims)
|
|
1116
1223
|
|
|
1117
1224
|
|
|
1118
1225
|
def squeeze(
|
|
@@ -1162,7 +1269,7 @@ def eye(
|
|
|
1162
1269
|
_validate_device(device)
|
|
1163
1270
|
n_cols = n_rows if n_cols is None else n_cols
|
|
1164
1271
|
dtype = jl_dtypes.float64 if dtype is None else dtype
|
|
1165
|
-
tns = jl.
|
|
1272
|
+
tns = jl.Galley.eye_python(n_rows, n_cols, k, dtype(False))
|
|
1166
1273
|
if format == "coo":
|
|
1167
1274
|
return Tensor(tns)
|
|
1168
1275
|
if format == "dense":
|
|
@@ -1316,16 +1423,48 @@ def logaddexp(x1: Tensor, x2: Tensor, /) -> Tensor:
|
|
|
1316
1423
|
return log(exp(x1) + exp(x2))
|
|
1317
1424
|
|
|
1318
1425
|
|
|
1426
|
+
def _truth(x: Tensor) -> Tensor:
|
|
1427
|
+
"""x as Bools, the truth the logical operators read (Galley.and and its
|
|
1428
|
+
relatives take only Bools)."""
|
|
1429
|
+
# eltype, unlike `dtype`, also reads lazy tensors (einop's operands).
|
|
1430
|
+
if isinstance(x, Tensor) and jl.eltype(x._obj) == jl.Bool:
|
|
1431
|
+
return x
|
|
1432
|
+
return x != 0
|
|
1433
|
+
|
|
1434
|
+
|
|
1435
|
+
def _logical(op: str, x1: Tensor, x2: Tensor) -> Tensor:
|
|
1436
|
+
if not isinstance(x1, Tensor):
|
|
1437
|
+
x1, x2 = x2, x1 # The logical operators commute.
|
|
1438
|
+
return _truth(x1)._elemwise_op(op, _truth(x2))
|
|
1439
|
+
|
|
1440
|
+
|
|
1319
1441
|
def logical_and(x1: Tensor, x2: Tensor, /) -> Tensor:
|
|
1320
|
-
return
|
|
1442
|
+
return _logical("Galley.and", x1, x2)
|
|
1321
1443
|
|
|
1322
1444
|
|
|
1323
1445
|
def logical_or(x1: Tensor, x2: Tensor, /) -> Tensor:
|
|
1324
|
-
return
|
|
1446
|
+
return _logical("Galley.or", x1, x2)
|
|
1325
1447
|
|
|
1326
1448
|
|
|
1327
1449
|
def logical_xor(x1: Tensor, x2: Tensor, /) -> Tensor:
|
|
1328
|
-
return
|
|
1450
|
+
return _logical("Galley.xor", x1, x2)
|
|
1451
|
+
|
|
1452
|
+
|
|
1453
|
+
def maximum(x1: Tensor, x2: Tensor, /) -> Tensor:
|
|
1454
|
+
if not isinstance(x1, Tensor):
|
|
1455
|
+
x1, x2 = x2, x1 # max commutes.
|
|
1456
|
+
return x1._elemwise_op("max", x2)
|
|
1457
|
+
|
|
1458
|
+
|
|
1459
|
+
def minimum(x1: Tensor, x2: Tensor, /) -> Tensor:
|
|
1460
|
+
if not isinstance(x1, Tensor):
|
|
1461
|
+
x1, x2 = x2, x1 # min commutes.
|
|
1462
|
+
return x1._elemwise_op("min", x2)
|
|
1463
|
+
|
|
1464
|
+
|
|
1465
|
+
def logical_not(x: Tensor, /) -> Tensor:
|
|
1466
|
+
# einop's `not` and `logical_not` resolve to this function.
|
|
1467
|
+
return x._elemwise_op("iszero")
|
|
1329
1468
|
|
|
1330
1469
|
|
|
1331
1470
|
def power(x1: Tensor, x2: Tensor, /) -> Tensor:
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|