galley-jl-python 0.1.0__tar.gz → 0.1.1__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.1}/PKG-INFO +6 -6
- {galley_jl_python-0.1.0 → galley_jl_python-0.1.1}/README.md +5 -5
- {galley_jl_python-0.1.0 → galley_jl_python-0.1.1}/pyproject.toml +1 -1
- {galley_jl_python-0.1.0 → galley_jl_python-0.1.1}/src/galley_jl_python/__init__.py +6 -0
- {galley_jl_python-0.1.0 → galley_jl_python-0.1.1}/src/galley_jl_python/compiled.py +7 -7
- {galley_jl_python-0.1.0 → galley_jl_python-0.1.1}/src/galley_jl_python/julia.py +4 -5
- {galley_jl_python-0.1.0 → galley_jl_python-0.1.1}/src/galley_jl_python/juliapkg.json +6 -3
- {galley_jl_python-0.1.0 → galley_jl_python-0.1.1}/src/galley_jl_python/juliapkg_dev.json +3 -3
- {galley_jl_python-0.1.0 → galley_jl_python-0.1.1}/src/galley_jl_python/linalg/_linalg.py +1 -1
- {galley_jl_python-0.1.0 → galley_jl_python-0.1.1}/src/galley_jl_python/tensor.py +155 -42
- {galley_jl_python-0.1.0 → galley_jl_python-0.1.1}/LICENSE +0 -0
- {galley_jl_python-0.1.0 → galley_jl_python-0.1.1}/src/galley_jl_python/_array_api_info.py +0 -0
- {galley_jl_python-0.1.0 → galley_jl_python-0.1.1}/src/galley_jl_python/_sysimage.py +0 -0
- {galley_jl_python-0.1.0 → galley_jl_python-0.1.1}/src/galley_jl_python/dtypes.py +0 -0
- {galley_jl_python-0.1.0 → galley_jl_python-0.1.1}/src/galley_jl_python/einstein.py +0 -0
- {galley_jl_python-0.1.0 → galley_jl_python-0.1.1}/src/galley_jl_python/errors.py +0 -0
- {galley_jl_python-0.1.0 → galley_jl_python-0.1.1}/src/galley_jl_python/finch_fused.py +0 -0
- {galley_jl_python-0.1.0 → galley_jl_python-0.1.1}/src/galley_jl_python/io.py +0 -0
- {galley_jl_python-0.1.0 → galley_jl_python-0.1.1}/src/galley_jl_python/levels.py +0 -0
- {galley_jl_python-0.1.0 → galley_jl_python-0.1.1}/src/galley_jl_python/linalg/__init__.py +0 -0
- {galley_jl_python-0.1.0 → galley_jl_python-0.1.1}/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.1
|
|
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.1"
|
|
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,53 @@ 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 _element_fill(dtype, fill_value):
|
|
72
|
+
"""fill_value as an Element level's fill for values of dtype.
|
|
73
|
+
|
|
74
|
+
Finch requires an isbits fill. Julia converts NumPy's numeric scalars to
|
|
75
|
+
Julia numbers, but not `np.bool_` or 0-d arrays (which `sparse.COO` uses for
|
|
76
|
+
its fill), so both are normalized first, and Booleans become Python bools.
|
|
77
|
+
"""
|
|
78
|
+
fill_value = np.asarray(fill_value).item()
|
|
79
|
+
if np.dtype(dtype) == np.bool_:
|
|
80
|
+
return bool(fill_value)
|
|
81
|
+
return np.dtype(dtype).type(fill_value)
|
|
82
|
+
|
|
83
|
+
|
|
40
84
|
class Tensor(_Display, SparseArray):
|
|
41
85
|
"""
|
|
42
|
-
A wrapper class for
|
|
86
|
+
A wrapper class for Galley.Tensor and Galley.SwizzleArray.
|
|
43
87
|
|
|
44
88
|
Constructors
|
|
45
89
|
------------
|
|
@@ -58,7 +102,7 @@ class Tensor(_Display, SparseArray):
|
|
|
58
102
|
|
|
59
103
|
Parameters
|
|
60
104
|
----------
|
|
61
|
-
obj : np.ndarray or scipy.sparse or Storage or
|
|
105
|
+
obj : np.ndarray or scipy.sparse or Storage or Galley.SwizzleArray
|
|
62
106
|
Input to construct a Tensor. It's a no-copy operation of for NumPy and
|
|
63
107
|
SciPy input. For Storage it's levels' description with order. The order
|
|
64
108
|
numbers the dimensions from the fastest to slowest. The leaf nodes have
|
|
@@ -75,7 +119,7 @@ class Tensor(_Display, SparseArray):
|
|
|
75
119
|
Returns
|
|
76
120
|
-------
|
|
77
121
|
Tensor
|
|
78
|
-
Python wrapper for
|
|
122
|
+
Python wrapper for a Galley Tensor.
|
|
79
123
|
|
|
80
124
|
Examples
|
|
81
125
|
--------
|
|
@@ -108,10 +152,10 @@ class Tensor(_Display, SparseArray):
|
|
|
108
152
|
fill_value: np.number | None = None,
|
|
109
153
|
copy: bool | None = None,
|
|
110
154
|
):
|
|
111
|
-
if isinstance(obj, int | float | complex | bool | list):
|
|
155
|
+
if isinstance(obj, int | float | complex | bool | list | tuple):
|
|
112
156
|
if copy is False:
|
|
113
157
|
raise ValueError(
|
|
114
|
-
"copy=False isn't supported for scalar inputs and Python
|
|
158
|
+
"copy=False isn't supported for scalar inputs and Python sequences"
|
|
115
159
|
)
|
|
116
160
|
obj = np.asarray(obj)
|
|
117
161
|
if fill_value is None:
|
|
@@ -130,11 +174,11 @@ class Tensor(_Display, SparseArray):
|
|
|
130
174
|
obj.order, self.get_lvl_ndim(obj.levels_descr._obj)
|
|
131
175
|
)
|
|
132
176
|
self._obj = jl.swizzle(jl.Tensor(obj.levels_descr._obj), *order)
|
|
133
|
-
elif jl.isa(obj, jl.
|
|
177
|
+
elif jl.isa(obj, jl.Galley.Tensor): # raw-Julia-object constructors
|
|
134
178
|
if copy:
|
|
135
179
|
self._raise_julia_copy_not_supported()
|
|
136
180
|
self._obj = jl.swizzle(obj, *tuple(range(1, jl.ndims(obj) + 1)))
|
|
137
|
-
elif jl.isa(obj, jl.
|
|
181
|
+
elif jl.isa(obj, jl.Galley.SwizzleArray) or jl.isa(obj, jl.Galley.LazyTensor):
|
|
138
182
|
if copy:
|
|
139
183
|
self._raise_julia_copy_not_supported()
|
|
140
184
|
self._obj = obj
|
|
@@ -165,14 +209,51 @@ class Tensor(_Display, SparseArray):
|
|
|
165
209
|
return self._elemwise_op("/", other)
|
|
166
210
|
|
|
167
211
|
def __floordiv__(self, other):
|
|
168
|
-
return self._elemwise_op("
|
|
212
|
+
return self._elemwise_op("Galley.fld_nothrow", other)
|
|
169
213
|
|
|
170
214
|
def __mod__(self, other):
|
|
171
|
-
return self._elemwise_op("
|
|
215
|
+
return self._elemwise_op("Galley.mod_nothrow", other)
|
|
172
216
|
|
|
173
217
|
def __pow__(self, other):
|
|
174
218
|
return self._elemwise_op("^", other)
|
|
175
219
|
|
|
220
|
+
# Reflected operators, for a scalar on the left (`1 - x`, `2 ** x`).
|
|
221
|
+
def __radd__(self, other):
|
|
222
|
+
return self._elemwise_op("+", other, reflected=True)
|
|
223
|
+
|
|
224
|
+
def __rsub__(self, other):
|
|
225
|
+
return self._elemwise_op("-", other, reflected=True)
|
|
226
|
+
|
|
227
|
+
def __rmul__(self, other):
|
|
228
|
+
return self._elemwise_op("*", other, reflected=True)
|
|
229
|
+
|
|
230
|
+
def __rtruediv__(self, other):
|
|
231
|
+
return self._elemwise_op("/", other, reflected=True)
|
|
232
|
+
|
|
233
|
+
def __rfloordiv__(self, other):
|
|
234
|
+
return self._elemwise_op("Galley.fld_nothrow", other, reflected=True)
|
|
235
|
+
|
|
236
|
+
def __rmod__(self, other):
|
|
237
|
+
return self._elemwise_op("Galley.mod_nothrow", other, reflected=True)
|
|
238
|
+
|
|
239
|
+
def __rpow__(self, other):
|
|
240
|
+
return self._elemwise_op("^", other, reflected=True)
|
|
241
|
+
|
|
242
|
+
def __rand__(self, other):
|
|
243
|
+
return self._elemwise_op("&", other, reflected=True)
|
|
244
|
+
|
|
245
|
+
def __ror__(self, other):
|
|
246
|
+
return self._elemwise_op("|", other, reflected=True)
|
|
247
|
+
|
|
248
|
+
def __rxor__(self, other):
|
|
249
|
+
return self._elemwise_op("xor", other, reflected=True)
|
|
250
|
+
|
|
251
|
+
def __rlshift__(self, other):
|
|
252
|
+
return self._elemwise_op("<<", other, reflected=True)
|
|
253
|
+
|
|
254
|
+
def __rrshift__(self, other):
|
|
255
|
+
return self._elemwise_op(">>", other, reflected=True)
|
|
256
|
+
|
|
176
257
|
@compiled()
|
|
177
258
|
def __matmul__(self, other: Tensor) -> Tensor:
|
|
178
259
|
if self.ndim == 0 or other.ndim == 0:
|
|
@@ -227,20 +308,20 @@ class Tensor(_Display, SparseArray):
|
|
|
227
308
|
def __ne__(self, other):
|
|
228
309
|
return self._elemwise_op("!=", other)
|
|
229
310
|
|
|
230
|
-
def _elemwise_op(
|
|
311
|
+
def _elemwise_op(
|
|
312
|
+
self, op: str, other: Tensor | None = None, *, reflected: bool = False
|
|
313
|
+
) -> Tensor:
|
|
231
314
|
if other is None:
|
|
232
315
|
result = jl.broadcast(jl.seval(op), self._obj)
|
|
233
316
|
else:
|
|
234
317
|
if np.isscalar(other):
|
|
235
|
-
other =
|
|
318
|
+
other = _weak_scalar(jl.eltype(self._obj), other)
|
|
236
319
|
else:
|
|
237
320
|
other = jl.permutedims(other._obj, tuple(range(other.ndim, 0, -1)))
|
|
238
321
|
# inverse swizzle, so `broadcast` appends new dims to the front
|
|
239
|
-
|
|
240
|
-
|
|
241
|
-
|
|
242
|
-
other,
|
|
243
|
-
)
|
|
322
|
+
this = jl.permutedims(self._obj, tuple(range(self.ndim, 0, -1)))
|
|
323
|
+
args = (other, this) if reflected else (this, other)
|
|
324
|
+
result = jl.broadcast(jl.seval(op), *args)
|
|
244
325
|
# swizzle back to the original order
|
|
245
326
|
result = jl.permutedims(result, tuple(range(jl.ndims(result), 0, -1)))
|
|
246
327
|
|
|
@@ -280,9 +361,11 @@ class Tensor(_Display, SparseArray):
|
|
|
280
361
|
key = _add_plus_one(key, self.shape)
|
|
281
362
|
|
|
282
363
|
result = self._obj[key]
|
|
283
|
-
if jl.isa(result, jl.
|
|
364
|
+
if jl.isa(result, jl.Galley.SwizzleArray) or jl.isa(
|
|
365
|
+
result, jl.Galley.LazyTensor
|
|
366
|
+
):
|
|
284
367
|
return Tensor(result)
|
|
285
|
-
if jl.isa(result, jl.
|
|
368
|
+
if jl.isa(result, jl.Galley.Tensor):
|
|
286
369
|
return Tensor(jl.swizzle(result, *range(1, jl.ndims(result) + 1)))
|
|
287
370
|
# the Array API requires a 0-d array, not a scalar, for integer-only keys
|
|
288
371
|
return Tensor(jl.swizzle(_scalar_to_0d(self.dtype, self.fill_value, result)))
|
|
@@ -311,7 +394,7 @@ class Tensor(_Display, SparseArray):
|
|
|
311
394
|
def _is_dense(self) -> bool:
|
|
312
395
|
lvl = self._obj.body.lvl
|
|
313
396
|
for _ in self.shape:
|
|
314
|
-
if not jl.isa(lvl, jl.
|
|
397
|
+
if not jl.isa(lvl, jl.Galley.Dense):
|
|
315
398
|
return False
|
|
316
399
|
lvl = lvl.lvl
|
|
317
400
|
return True
|
|
@@ -320,6 +403,12 @@ class Tensor(_Display, SparseArray):
|
|
|
320
403
|
def _order(self) -> tuple[int, ...]:
|
|
321
404
|
return jl.typeof(self._obj).parameters[1]
|
|
322
405
|
|
|
406
|
+
@property
|
|
407
|
+
def T(self) -> Tensor:
|
|
408
|
+
if self.ndim != 2:
|
|
409
|
+
raise ValueError(f"`T` requires a 2-dimensional array, got {self.ndim=}.")
|
|
410
|
+
return self.permute_dims((1, 0))
|
|
411
|
+
|
|
323
412
|
@property
|
|
324
413
|
def mT(self) -> Tensor:
|
|
325
414
|
axes = list(range(self.ndim))
|
|
@@ -340,7 +429,7 @@ class Tensor(_Display, SparseArray):
|
|
|
340
429
|
return self
|
|
341
430
|
|
|
342
431
|
def is_computed(self) -> bool:
|
|
343
|
-
return not jl.isa(self._obj, jl.
|
|
432
|
+
return not jl.isa(self._obj, jl.Galley.LazyTensor)
|
|
344
433
|
|
|
345
434
|
@classmethod
|
|
346
435
|
def preprocess_order(cls, order: OrderType, ndim: int) -> tuple[int, ...]:
|
|
@@ -370,7 +459,7 @@ class Tensor(_Display, SparseArray):
|
|
|
370
459
|
while True:
|
|
371
460
|
ndim += 1
|
|
372
461
|
lvl = lvl.lvl
|
|
373
|
-
if jl.isa(lvl, jl.
|
|
462
|
+
if jl.isa(lvl, jl.Galley.Element):
|
|
374
463
|
break
|
|
375
464
|
return ndim
|
|
376
465
|
|
|
@@ -433,12 +522,7 @@ class Tensor(_Display, SparseArray):
|
|
|
433
522
|
order = cls.preprocess_order(order_char, arr.ndim)
|
|
434
523
|
inv_order = tuple(i - 1 for i in jl.invperm(order))
|
|
435
524
|
|
|
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)
|
|
525
|
+
fill_value = _element_fill(arr.dtype, fill_value)
|
|
442
526
|
lvl = Element(fill_value, arr.reshape(-1, order=order_char))
|
|
443
527
|
for i in inv_order:
|
|
444
528
|
lvl = Dense(lvl, arr.shape[i])
|
|
@@ -503,11 +587,10 @@ class Tensor(_Display, SparseArray):
|
|
|
503
587
|
def construct_coo_jl_object(
|
|
504
588
|
cls, coords, data, shape, order, fill_value=0.0
|
|
505
589
|
) -> JuliaObj:
|
|
506
|
-
assert len(coords) == 2
|
|
507
590
|
ndim = len(shape)
|
|
508
591
|
order = cls.preprocess_order(order, ndim)
|
|
509
592
|
|
|
510
|
-
lvl = jl.Element(data.dtype
|
|
593
|
+
lvl = jl.Element(_element_fill(data.dtype, fill_value), data)
|
|
511
594
|
ptr = jl.Vector[jl.Int]([1, len(data) + 1])
|
|
512
595
|
tbl = tuple(jl.PlusOneVector(arr) for arr in coords)
|
|
513
596
|
|
|
@@ -532,11 +615,10 @@ class Tensor(_Display, SparseArray):
|
|
|
532
615
|
assert len(shape) == 2
|
|
533
616
|
|
|
534
617
|
data, indices, indptr = arg
|
|
535
|
-
dtype = data.dtype.type
|
|
536
618
|
indices = jl.PlusOneVector(indices)
|
|
537
619
|
indptr = jl.PlusOneVector(indptr)
|
|
538
620
|
|
|
539
|
-
lvl = jl.Element(dtype
|
|
621
|
+
lvl = jl.Element(_element_fill(data.dtype, fill_value), data)
|
|
540
622
|
return jl.swizzle(
|
|
541
623
|
jl.Tensor(
|
|
542
624
|
jl.Dense(jl.SparseList(lvl, shape[0], indptr, indices), shape[1])
|
|
@@ -579,7 +661,6 @@ class Tensor(_Display, SparseArray):
|
|
|
579
661
|
assert isinstance(arg, tuple) and len(arg) == 3
|
|
580
662
|
|
|
581
663
|
data, indices_list, indptr_list = arg
|
|
582
|
-
dtype = data.dtype.type
|
|
583
664
|
|
|
584
665
|
assert len(indices_list) == len(shape) - 1
|
|
585
666
|
assert len(indptr_list) == len(shape) - 1
|
|
@@ -587,7 +668,7 @@ class Tensor(_Display, SparseArray):
|
|
|
587
668
|
indices_list = [jl.PlusOneVector(i) for i in indices_list]
|
|
588
669
|
indptr_list = [jl.PlusOneVector(i) for i in indptr_list]
|
|
589
670
|
|
|
590
|
-
lvl = jl.Element(dtype
|
|
671
|
+
lvl = jl.Element(_element_fill(data.dtype, fill_value), data)
|
|
591
672
|
for size, indices, indptr in zip(
|
|
592
673
|
shape[:-1], indices_list, indptr_list, strict=False
|
|
593
674
|
):
|
|
@@ -965,7 +1046,7 @@ def _reduce_sum_prod(
|
|
|
965
1046
|
|
|
966
1047
|
result = Tensor(result)
|
|
967
1048
|
|
|
968
|
-
if jl.isa(result._obj, jl.
|
|
1049
|
+
if jl.isa(result._obj, jl.Galley.LazyTensor):
|
|
969
1050
|
if dtype is not None:
|
|
970
1051
|
raise ValueError(
|
|
971
1052
|
"`dtype` keyword for `sum` and `prod` in the lazy mode isn't supported"
|
|
@@ -1102,7 +1183,7 @@ def argmin(
|
|
|
1102
1183
|
axis: int | None = None,
|
|
1103
1184
|
keepdims: bool = False,
|
|
1104
1185
|
) -> Tensor:
|
|
1105
|
-
return _reduce(x, jl.
|
|
1186
|
+
return _reduce(x, jl.Galley.argmin_python, axis, keepdims)
|
|
1106
1187
|
|
|
1107
1188
|
|
|
1108
1189
|
def argmax(
|
|
@@ -1112,7 +1193,7 @@ def argmax(
|
|
|
1112
1193
|
axis: int | None = None,
|
|
1113
1194
|
keepdims: bool = False,
|
|
1114
1195
|
) -> Tensor:
|
|
1115
|
-
return _reduce(x, jl.
|
|
1196
|
+
return _reduce(x, jl.Galley.argmax_python, axis, keepdims)
|
|
1116
1197
|
|
|
1117
1198
|
|
|
1118
1199
|
def squeeze(
|
|
@@ -1162,7 +1243,7 @@ def eye(
|
|
|
1162
1243
|
_validate_device(device)
|
|
1163
1244
|
n_cols = n_rows if n_cols is None else n_cols
|
|
1164
1245
|
dtype = jl_dtypes.float64 if dtype is None else dtype
|
|
1165
|
-
tns = jl.
|
|
1246
|
+
tns = jl.Galley.eye_python(n_rows, n_cols, k, dtype(False))
|
|
1166
1247
|
if format == "coo":
|
|
1167
1248
|
return Tensor(tns)
|
|
1168
1249
|
if format == "dense":
|
|
@@ -1316,16 +1397,48 @@ def logaddexp(x1: Tensor, x2: Tensor, /) -> Tensor:
|
|
|
1316
1397
|
return log(exp(x1) + exp(x2))
|
|
1317
1398
|
|
|
1318
1399
|
|
|
1400
|
+
def _truth(x: Tensor) -> Tensor:
|
|
1401
|
+
"""x as Bools, the truth the logical operators read (Galley.and and its
|
|
1402
|
+
relatives take only Bools)."""
|
|
1403
|
+
# eltype, unlike `dtype`, also reads lazy tensors (einop's operands).
|
|
1404
|
+
if isinstance(x, Tensor) and jl.eltype(x._obj) == jl.Bool:
|
|
1405
|
+
return x
|
|
1406
|
+
return x != 0
|
|
1407
|
+
|
|
1408
|
+
|
|
1409
|
+
def _logical(op: str, x1: Tensor, x2: Tensor) -> Tensor:
|
|
1410
|
+
if not isinstance(x1, Tensor):
|
|
1411
|
+
x1, x2 = x2, x1 # The logical operators commute.
|
|
1412
|
+
return _truth(x1)._elemwise_op(op, _truth(x2))
|
|
1413
|
+
|
|
1414
|
+
|
|
1319
1415
|
def logical_and(x1: Tensor, x2: Tensor, /) -> Tensor:
|
|
1320
|
-
return
|
|
1416
|
+
return _logical("Galley.and", x1, x2)
|
|
1321
1417
|
|
|
1322
1418
|
|
|
1323
1419
|
def logical_or(x1: Tensor, x2: Tensor, /) -> Tensor:
|
|
1324
|
-
return
|
|
1420
|
+
return _logical("Galley.or", x1, x2)
|
|
1325
1421
|
|
|
1326
1422
|
|
|
1327
1423
|
def logical_xor(x1: Tensor, x2: Tensor, /) -> Tensor:
|
|
1328
|
-
return
|
|
1424
|
+
return _logical("Galley.xor", x1, x2)
|
|
1425
|
+
|
|
1426
|
+
|
|
1427
|
+
def maximum(x1: Tensor, x2: Tensor, /) -> Tensor:
|
|
1428
|
+
if not isinstance(x1, Tensor):
|
|
1429
|
+
x1, x2 = x2, x1 # max commutes.
|
|
1430
|
+
return x1._elemwise_op("max", x2)
|
|
1431
|
+
|
|
1432
|
+
|
|
1433
|
+
def minimum(x1: Tensor, x2: Tensor, /) -> Tensor:
|
|
1434
|
+
if not isinstance(x1, Tensor):
|
|
1435
|
+
x1, x2 = x2, x1 # min commutes.
|
|
1436
|
+
return x1._elemwise_op("min", x2)
|
|
1437
|
+
|
|
1438
|
+
|
|
1439
|
+
def logical_not(x: Tensor, /) -> Tensor:
|
|
1440
|
+
# einop's `not` and `logical_not` resolve to this function.
|
|
1441
|
+
return x._elemwise_op("iszero")
|
|
1329
1442
|
|
|
1330
1443
|
|
|
1331
1444
|
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
|