brainstate 0.2.0__py2.py3-none-any.whl → 0.2.1__py2.py3-none-any.whl

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.
Files changed (112) hide show
  1. brainstate/__init__.py +169 -169
  2. brainstate/_compatible_import.py +340 -340
  3. brainstate/_compatible_import_test.py +681 -681
  4. brainstate/_deprecation.py +210 -210
  5. brainstate/_deprecation_test.py +2319 -2319
  6. brainstate/_error.py +45 -45
  7. brainstate/_state.py +1652 -1652
  8. brainstate/_state_test.py +52 -52
  9. brainstate/_utils.py +47 -47
  10. brainstate/environ.py +1495 -1495
  11. brainstate/environ_test.py +1223 -1223
  12. brainstate/graph/__init__.py +22 -22
  13. brainstate/graph/_node.py +240 -240
  14. brainstate/graph/_node_test.py +589 -589
  15. brainstate/graph/_operation.py +1624 -1624
  16. brainstate/graph/_operation_test.py +1147 -1147
  17. brainstate/mixin.py +1433 -1433
  18. brainstate/mixin_test.py +1017 -1017
  19. brainstate/nn/__init__.py +137 -137
  20. brainstate/nn/_activations.py +1100 -1100
  21. brainstate/nn/_activations_test.py +354 -354
  22. brainstate/nn/_collective_ops.py +633 -633
  23. brainstate/nn/_collective_ops_test.py +774 -774
  24. brainstate/nn/_common.py +226 -226
  25. brainstate/nn/_common_test.py +154 -154
  26. brainstate/nn/_conv.py +2010 -2010
  27. brainstate/nn/_conv_test.py +849 -849
  28. brainstate/nn/_delay.py +575 -575
  29. brainstate/nn/_delay_test.py +243 -243
  30. brainstate/nn/_dropout.py +618 -618
  31. brainstate/nn/_dropout_test.py +477 -477
  32. brainstate/nn/_dynamics.py +1267 -1267
  33. brainstate/nn/_dynamics_test.py +67 -67
  34. brainstate/nn/_elementwise.py +1298 -1298
  35. brainstate/nn/_elementwise_test.py +829 -829
  36. brainstate/nn/_embedding.py +408 -408
  37. brainstate/nn/_embedding_test.py +156 -156
  38. brainstate/nn/_event_fixedprob.py +233 -233
  39. brainstate/nn/_event_fixedprob_test.py +115 -115
  40. brainstate/nn/_event_linear.py +83 -83
  41. brainstate/nn/_event_linear_test.py +121 -121
  42. brainstate/nn/_exp_euler.py +254 -254
  43. brainstate/nn/_exp_euler_test.py +377 -377
  44. brainstate/nn/_linear.py +744 -744
  45. brainstate/nn/_linear_test.py +475 -475
  46. brainstate/nn/_metrics.py +1070 -1070
  47. brainstate/nn/_metrics_test.py +611 -611
  48. brainstate/nn/_module.py +384 -384
  49. brainstate/nn/_module_test.py +40 -40
  50. brainstate/nn/_normalizations.py +1334 -1334
  51. brainstate/nn/_normalizations_test.py +699 -699
  52. brainstate/nn/_paddings.py +1020 -1020
  53. brainstate/nn/_paddings_test.py +722 -722
  54. brainstate/nn/_poolings.py +2239 -2239
  55. brainstate/nn/_poolings_test.py +952 -952
  56. brainstate/nn/_rnns.py +946 -946
  57. brainstate/nn/_rnns_test.py +592 -592
  58. brainstate/nn/_utils.py +216 -216
  59. brainstate/nn/_utils_test.py +401 -401
  60. brainstate/nn/init.py +809 -809
  61. brainstate/nn/init_test.py +180 -180
  62. brainstate/random/__init__.py +270 -270
  63. brainstate/random/_rand_funs.py +3938 -3938
  64. brainstate/random/_rand_funs_test.py +640 -640
  65. brainstate/random/_rand_seed.py +675 -675
  66. brainstate/random/_rand_seed_test.py +48 -48
  67. brainstate/random/_rand_state.py +1617 -1617
  68. brainstate/random/_rand_state_test.py +551 -551
  69. brainstate/transform/__init__.py +59 -59
  70. brainstate/transform/_ad_checkpoint.py +176 -176
  71. brainstate/transform/_ad_checkpoint_test.py +49 -49
  72. brainstate/transform/_autograd.py +1025 -1025
  73. brainstate/transform/_autograd_test.py +1289 -1289
  74. brainstate/transform/_conditions.py +316 -316
  75. brainstate/transform/_conditions_test.py +220 -220
  76. brainstate/transform/_error_if.py +94 -94
  77. brainstate/transform/_error_if_test.py +52 -52
  78. brainstate/transform/_eval_shape.py +145 -145
  79. brainstate/transform/_eval_shape_test.py +38 -38
  80. brainstate/transform/_jit.py +399 -399
  81. brainstate/transform/_jit_test.py +143 -143
  82. brainstate/transform/_loop_collect_return.py +675 -675
  83. brainstate/transform/_loop_collect_return_test.py +58 -58
  84. brainstate/transform/_loop_no_collection.py +283 -283
  85. brainstate/transform/_loop_no_collection_test.py +50 -50
  86. brainstate/transform/_make_jaxpr.py +2016 -2016
  87. brainstate/transform/_make_jaxpr_test.py +1510 -1510
  88. brainstate/transform/_mapping.py +529 -529
  89. brainstate/transform/_mapping_test.py +194 -194
  90. brainstate/transform/_progress_bar.py +255 -255
  91. brainstate/transform/_random.py +171 -171
  92. brainstate/transform/_unvmap.py +256 -256
  93. brainstate/transform/_util.py +286 -286
  94. brainstate/typing.py +837 -837
  95. brainstate/typing_test.py +780 -780
  96. brainstate/util/__init__.py +27 -27
  97. brainstate/util/_others.py +1024 -1024
  98. brainstate/util/_others_test.py +962 -962
  99. brainstate/util/_pretty_pytree.py +1301 -1301
  100. brainstate/util/_pretty_pytree_test.py +675 -675
  101. brainstate/util/_pretty_repr.py +462 -462
  102. brainstate/util/_pretty_repr_test.py +696 -696
  103. brainstate/util/filter.py +945 -945
  104. brainstate/util/filter_test.py +911 -911
  105. brainstate/util/struct.py +910 -910
  106. brainstate/util/struct_test.py +602 -602
  107. {brainstate-0.2.0.dist-info → brainstate-0.2.1.dist-info}/METADATA +108 -108
  108. brainstate-0.2.1.dist-info/RECORD +111 -0
  109. {brainstate-0.2.0.dist-info → brainstate-0.2.1.dist-info}/licenses/LICENSE +202 -202
  110. brainstate-0.2.0.dist-info/RECORD +0 -111
  111. {brainstate-0.2.0.dist-info → brainstate-0.2.1.dist-info}/WHEEL +0 -0
  112. {brainstate-0.2.0.dist-info → brainstate-0.2.1.dist-info}/top_level.txt +0 -0
@@ -1,255 +1,255 @@
1
- # Copyright 2024 BrainX Ecosystem Limited. All Rights Reserved.
2
- #
3
- # Licensed under the Apache License, Version 2.0 (the "License");
4
- # you may not use this file except in compliance with the License.
5
- # You may obtain a copy of the License at
6
- #
7
- # http://www.apache.org/licenses/LICENSE-2.0
8
- #
9
- # Unless required by applicable law or agreed to in writing, software
10
- # distributed under the License is distributed on an "AS IS" BASIS,
11
- # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
- # See the License for the specific language governing permissions and
13
- # limitations under the License.
14
- # ==============================================================================
15
-
16
-
17
- import copy
18
- import importlib.util
19
- from typing import Optional, Callable, Any, Tuple, Dict
20
-
21
- import jax
22
-
23
- tqdm_installed = importlib.util.find_spec('tqdm') is not None
24
-
25
- __all__ = [
26
- 'ProgressBar',
27
- ]
28
-
29
- Index = int
30
- Carray = Any
31
- Output = Any
32
-
33
-
34
- class ProgressBar(object):
35
- """
36
- A progress bar for tracking the progress of a jitted for-loop computation.
37
-
38
- It can be used in :py:func:`for_loop`, :py:func:`checkpointed_for_loop`, :py:func:`scan`,
39
- and :py:func:`checkpointed_scan` functions. Or any other jitted function that uses
40
- a for-loop.
41
-
42
- The message displayed in the progress bar can be customized by the following two methods:
43
-
44
- 1. By passing a string to the `desc` argument.
45
- 2. By passing a tuple with a string and a callable function to the `desc` argument. The callable
46
- function should take a dictionary as input and return a dictionary. The returned dictionary
47
- will be used to format the string.
48
-
49
- In the second case, ``"i"`` denotes the iteration number and other keys can be computed from the
50
- loop outputs and carry values.
51
-
52
- Parameters
53
- ----------
54
- freq : int, optional
55
- The frequency at which to print the progress bar. If not specified, the progress
56
- bar will be printed every 5% of the total iterations.
57
- count : int, optional
58
- The number of times to print the progress bar. If not specified, the progress
59
- bar will be printed every 5% of the total iterations. Cannot be used together with `freq`.
60
- desc : str or tuple, optional
61
- A description of the progress bar. If not specified, a default message will be
62
- displayed. Can be either a string or a tuple of (format_string, format_function).
63
- **kwargs
64
- Additional keyword arguments to pass to the progress bar.
65
-
66
- Examples
67
- --------
68
- Basic usage with default description:
69
-
70
- .. code-block:: python
71
-
72
- >>> import brainstate
73
- >>> import jax.numpy as jnp
74
- >>>
75
- >>> def loop_fn(x):
76
- ... return x ** 2
77
- >>>
78
- >>> xs = jnp.arange(100)
79
- >>> pbar = brainstate.transform.ProgressBar()
80
- >>> results = brainstate.transform.for_loop(loop_fn, xs, pbar=pbar)
81
-
82
- With custom description string:
83
-
84
- .. code-block:: python
85
-
86
- >>> pbar = brainstate.transform.ProgressBar(desc="Running 1000 iterations")
87
- >>> results = brainstate.transform.for_loop(loop_fn, xs, pbar=pbar)
88
-
89
- With frequency control:
90
-
91
- .. code-block:: python
92
-
93
- >>> # Update every 10 iterations
94
- >>> pbar = brainstate.transform.ProgressBar(freq=10)
95
- >>> results = brainstate.transform.for_loop(loop_fn, xs, pbar=pbar)
96
- >>>
97
- >>> # Update exactly 20 times during execution
98
- >>> pbar = brainstate.transform.ProgressBar(count=20)
99
- >>> results = brainstate.transform.for_loop(loop_fn, xs, pbar=pbar)
100
-
101
- With dynamic description based on loop variables:
102
-
103
- .. code-block:: python
104
-
105
- >>> state = brainstate.State(1.0)
106
- >>>
107
- >>> def loop_fn(x):
108
- ... state.value += x
109
- ... loss = jnp.sum(x ** 2)
110
- ... return loss
111
- >>>
112
- >>> def format_desc(data):
113
- ... return {"i": data["i"], "loss": data["y"], "state": data["carry"]}
114
- >>>
115
- >>> pbar = brainstate.transform.ProgressBar(
116
- ... desc=("Iteration {i}, loss = {loss:.4f}, state = {state:.2f}", format_desc)
117
- ... )
118
- >>> results = brainstate.transform.for_loop(loop_fn, xs, pbar=pbar)
119
-
120
- With scan function:
121
-
122
- .. code-block:: python
123
-
124
- >>> def scan_fn(carry, x):
125
- ... new_carry = carry + x
126
- ... return new_carry, new_carry ** 2
127
- >>>
128
- >>> init_carry = 0.0
129
- >>> pbar = brainstate.transform.ProgressBar(freq=5)
130
- >>> final_carry, ys = brainstate.transform.scan(scan_fn, init_carry, xs, pbar=pbar)
131
- """
132
- __module__ = "brainstate.transform"
133
-
134
- def __init__(
135
- self,
136
- freq: Optional[int] = None,
137
- count: Optional[int] = None,
138
- desc: Optional[Tuple[str, Callable[[Dict], Dict]] | str] = None,
139
- **kwargs
140
- ):
141
- # print rate
142
- self.print_freq = freq
143
- if isinstance(freq, int):
144
- assert freq > 0, "Print rate should be > 0."
145
-
146
- # print count
147
- self.print_count = count
148
- if self.print_freq is not None and self.print_count is not None:
149
- raise ValueError("Cannot specify both count and freq.")
150
-
151
- # other parameters
152
- for kwarg in ("total", "mininterval", "maxinterval", "miniters"):
153
- kwargs.pop(kwarg, None)
154
- self.kwargs = kwargs
155
-
156
- # description
157
- if desc is not None:
158
- if isinstance(desc, str):
159
- pass
160
- else:
161
- assert isinstance(desc, (tuple, list)), 'Description should be a tuple or list.'
162
- assert isinstance(desc[0], str), 'Description should be a string.'
163
- assert callable(desc[1]), 'Description should be a callable.'
164
- self.desc = desc
165
-
166
- # check if tqdm is installed
167
- if not tqdm_installed:
168
- raise ImportError("tqdm is not installed.")
169
-
170
- def init(self, n: int):
171
- kwargs = copy.copy(self.kwargs)
172
- freq = self.print_freq
173
- count = self.print_count
174
- if count is not None:
175
- freq, remainder = divmod(n, count)
176
- if freq == 0:
177
- raise ValueError(f"Count {count} is too large for n {n}.")
178
- elif freq is None:
179
- if n > 20:
180
- freq = int(n / 20)
181
- else:
182
- freq = 1
183
- remainder = n % freq
184
- else:
185
- if freq < 1:
186
- raise ValueError(f"Print rate should be > 0 got {freq}")
187
- elif freq > n:
188
- raise ValueError("Print rate should be less than the "
189
- f"number of steps {n}, got {freq}")
190
- remainder = n % freq
191
-
192
- message = f"Running for {n:,} iterations" if self.desc is None else self.desc
193
- return ProgressBarRunner(n, freq, remainder, message, **kwargs)
194
-
195
-
196
- class ProgressBarRunner(object):
197
- __module__ = "brainstate.transform"
198
-
199
- def __init__(
200
- self,
201
- n: int,
202
- print_freq: int,
203
- remainder: int,
204
- message: str | Tuple[str, Callable[[Dict], Dict]],
205
- **kwargs
206
- ):
207
- self.tqdm_bars = {}
208
- self.kwargs = kwargs
209
- self.n = n
210
- self.print_freq = print_freq
211
- self.remainder = remainder
212
- self.message = message
213
-
214
- def _define_tqdm(self, x: dict):
215
- from tqdm.auto import tqdm
216
- self.tqdm_bars[0] = tqdm(range(self.n), **self.kwargs)
217
- if isinstance(self.message, str):
218
- self.tqdm_bars[0].set_description(self.message, refresh=False)
219
- else:
220
- self.tqdm_bars[0].set_description(self.message[0].format(**x), refresh=True)
221
-
222
- def _update_tqdm(self, x: dict):
223
- self.tqdm_bars[0].update(self.print_freq)
224
- if not isinstance(self.message, str):
225
- self.tqdm_bars[0].set_description(self.message[0].format(**x), refresh=True)
226
-
227
- def _close_tqdm(self, x: dict):
228
- if self.remainder > 0:
229
- self.tqdm_bars[0].update(self.remainder)
230
- if not isinstance(self.message, str):
231
- self.tqdm_bars[0].set_description(self.message[0].format(**x), refresh=True)
232
- self.tqdm_bars[0].close()
233
-
234
- def __call__(self, iter_num, **kwargs):
235
- data = dict() if isinstance(self.message, str) else self.message[1](dict(i=iter_num, **kwargs))
236
- assert isinstance(data, dict), 'Description function should return a dictionary.'
237
-
238
- _ = jax.lax.cond(
239
- iter_num == 0,
240
- lambda x: jax.debug.callback(self._define_tqdm, x, ordered=True),
241
- lambda x: None,
242
- data
243
- )
244
- _ = jax.lax.cond(
245
- iter_num % self.print_freq == (self.print_freq - 1),
246
- lambda x: jax.debug.callback(self._update_tqdm, x, ordered=True),
247
- lambda x: None,
248
- data
249
- )
250
- _ = jax.lax.cond(
251
- iter_num == self.n - 1,
252
- lambda x: jax.debug.callback(self._close_tqdm, x, ordered=True),
253
- lambda x: None,
254
- data
255
- )
1
+ # Copyright 2024 BrainX Ecosystem Limited. All Rights Reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+ # ==============================================================================
15
+
16
+
17
+ import copy
18
+ import importlib.util
19
+ from typing import Optional, Callable, Any, Tuple, Dict
20
+
21
+ import jax
22
+
23
+ tqdm_installed = importlib.util.find_spec('tqdm') is not None
24
+
25
+ __all__ = [
26
+ 'ProgressBar',
27
+ ]
28
+
29
+ Index = int
30
+ Carray = Any
31
+ Output = Any
32
+
33
+
34
+ class ProgressBar(object):
35
+ """
36
+ A progress bar for tracking the progress of a jitted for-loop computation.
37
+
38
+ It can be used in :py:func:`for_loop`, :py:func:`checkpointed_for_loop`, :py:func:`scan`,
39
+ and :py:func:`checkpointed_scan` functions. Or any other jitted function that uses
40
+ a for-loop.
41
+
42
+ The message displayed in the progress bar can be customized by the following two methods:
43
+
44
+ 1. By passing a string to the `desc` argument.
45
+ 2. By passing a tuple with a string and a callable function to the `desc` argument. The callable
46
+ function should take a dictionary as input and return a dictionary. The returned dictionary
47
+ will be used to format the string.
48
+
49
+ In the second case, ``"i"`` denotes the iteration number and other keys can be computed from the
50
+ loop outputs and carry values.
51
+
52
+ Parameters
53
+ ----------
54
+ freq : int, optional
55
+ The frequency at which to print the progress bar. If not specified, the progress
56
+ bar will be printed every 5% of the total iterations.
57
+ count : int, optional
58
+ The number of times to print the progress bar. If not specified, the progress
59
+ bar will be printed every 5% of the total iterations. Cannot be used together with `freq`.
60
+ desc : str or tuple, optional
61
+ A description of the progress bar. If not specified, a default message will be
62
+ displayed. Can be either a string or a tuple of (format_string, format_function).
63
+ **kwargs
64
+ Additional keyword arguments to pass to the progress bar.
65
+
66
+ Examples
67
+ --------
68
+ Basic usage with default description:
69
+
70
+ .. code-block:: python
71
+
72
+ >>> import brainstate
73
+ >>> import jax.numpy as jnp
74
+ >>>
75
+ >>> def loop_fn(x):
76
+ ... return x ** 2
77
+ >>>
78
+ >>> xs = jnp.arange(100)
79
+ >>> pbar = brainstate.transform.ProgressBar()
80
+ >>> results = brainstate.transform.for_loop(loop_fn, xs, pbar=pbar)
81
+
82
+ With custom description string:
83
+
84
+ .. code-block:: python
85
+
86
+ >>> pbar = brainstate.transform.ProgressBar(desc="Running 1000 iterations")
87
+ >>> results = brainstate.transform.for_loop(loop_fn, xs, pbar=pbar)
88
+
89
+ With frequency control:
90
+
91
+ .. code-block:: python
92
+
93
+ >>> # Update every 10 iterations
94
+ >>> pbar = brainstate.transform.ProgressBar(freq=10)
95
+ >>> results = brainstate.transform.for_loop(loop_fn, xs, pbar=pbar)
96
+ >>>
97
+ >>> # Update exactly 20 times during execution
98
+ >>> pbar = brainstate.transform.ProgressBar(count=20)
99
+ >>> results = brainstate.transform.for_loop(loop_fn, xs, pbar=pbar)
100
+
101
+ With dynamic description based on loop variables:
102
+
103
+ .. code-block:: python
104
+
105
+ >>> state = brainstate.State(1.0)
106
+ >>>
107
+ >>> def loop_fn(x):
108
+ ... state.value += x
109
+ ... loss = jnp.sum(x ** 2)
110
+ ... return loss
111
+ >>>
112
+ >>> def format_desc(data):
113
+ ... return {"i": data["i"], "loss": data["y"], "state": data["carry"]}
114
+ >>>
115
+ >>> pbar = brainstate.transform.ProgressBar(
116
+ ... desc=("Iteration {i}, loss = {loss:.4f}, state = {state:.2f}", format_desc)
117
+ ... )
118
+ >>> results = brainstate.transform.for_loop(loop_fn, xs, pbar=pbar)
119
+
120
+ With scan function:
121
+
122
+ .. code-block:: python
123
+
124
+ >>> def scan_fn(carry, x):
125
+ ... new_carry = carry + x
126
+ ... return new_carry, new_carry ** 2
127
+ >>>
128
+ >>> init_carry = 0.0
129
+ >>> pbar = brainstate.transform.ProgressBar(freq=5)
130
+ >>> final_carry, ys = brainstate.transform.scan(scan_fn, init_carry, xs, pbar=pbar)
131
+ """
132
+ __module__ = "brainstate.transform"
133
+
134
+ def __init__(
135
+ self,
136
+ freq: Optional[int] = None,
137
+ count: Optional[int] = None,
138
+ desc: Optional[Tuple[str, Callable[[Dict], Dict]] | str] = None,
139
+ **kwargs
140
+ ):
141
+ # print rate
142
+ self.print_freq = freq
143
+ if isinstance(freq, int):
144
+ assert freq > 0, "Print rate should be > 0."
145
+
146
+ # print count
147
+ self.print_count = count
148
+ if self.print_freq is not None and self.print_count is not None:
149
+ raise ValueError("Cannot specify both count and freq.")
150
+
151
+ # other parameters
152
+ for kwarg in ("total", "mininterval", "maxinterval", "miniters"):
153
+ kwargs.pop(kwarg, None)
154
+ self.kwargs = kwargs
155
+
156
+ # description
157
+ if desc is not None:
158
+ if isinstance(desc, str):
159
+ pass
160
+ else:
161
+ assert isinstance(desc, (tuple, list)), 'Description should be a tuple or list.'
162
+ assert isinstance(desc[0], str), 'Description should be a string.'
163
+ assert callable(desc[1]), 'Description should be a callable.'
164
+ self.desc = desc
165
+
166
+ # check if tqdm is installed
167
+ if not tqdm_installed:
168
+ raise ImportError("tqdm is not installed.")
169
+
170
+ def init(self, n: int):
171
+ kwargs = copy.copy(self.kwargs)
172
+ freq = self.print_freq
173
+ count = self.print_count
174
+ if count is not None:
175
+ freq, remainder = divmod(n, count)
176
+ if freq == 0:
177
+ raise ValueError(f"Count {count} is too large for n {n}.")
178
+ elif freq is None:
179
+ if n > 20:
180
+ freq = int(n / 20)
181
+ else:
182
+ freq = 1
183
+ remainder = n % freq
184
+ else:
185
+ if freq < 1:
186
+ raise ValueError(f"Print rate should be > 0 got {freq}")
187
+ elif freq > n:
188
+ raise ValueError("Print rate should be less than the "
189
+ f"number of steps {n}, got {freq}")
190
+ remainder = n % freq
191
+
192
+ message = f"Running for {n:,} iterations" if self.desc is None else self.desc
193
+ return ProgressBarRunner(n, freq, remainder, message, **kwargs)
194
+
195
+
196
+ class ProgressBarRunner(object):
197
+ __module__ = "brainstate.transform"
198
+
199
+ def __init__(
200
+ self,
201
+ n: int,
202
+ print_freq: int,
203
+ remainder: int,
204
+ message: str | Tuple[str, Callable[[Dict], Dict]],
205
+ **kwargs
206
+ ):
207
+ self.tqdm_bars = {}
208
+ self.kwargs = kwargs
209
+ self.n = n
210
+ self.print_freq = print_freq
211
+ self.remainder = remainder
212
+ self.message = message
213
+
214
+ def _define_tqdm(self, x: dict):
215
+ from tqdm.auto import tqdm
216
+ self.tqdm_bars[0] = tqdm(range(self.n), **self.kwargs)
217
+ if isinstance(self.message, str):
218
+ self.tqdm_bars[0].set_description(self.message, refresh=False)
219
+ else:
220
+ self.tqdm_bars[0].set_description(self.message[0].format(**x), refresh=True)
221
+
222
+ def _update_tqdm(self, x: dict):
223
+ self.tqdm_bars[0].update(self.print_freq)
224
+ if not isinstance(self.message, str):
225
+ self.tqdm_bars[0].set_description(self.message[0].format(**x), refresh=True)
226
+
227
+ def _close_tqdm(self, x: dict):
228
+ if self.remainder > 0:
229
+ self.tqdm_bars[0].update(self.remainder)
230
+ if not isinstance(self.message, str):
231
+ self.tqdm_bars[0].set_description(self.message[0].format(**x), refresh=True)
232
+ self.tqdm_bars[0].close()
233
+
234
+ def __call__(self, iter_num, **kwargs):
235
+ data = dict() if isinstance(self.message, str) else self.message[1](dict(i=iter_num, **kwargs))
236
+ assert isinstance(data, dict), 'Description function should return a dictionary.'
237
+
238
+ _ = jax.lax.cond(
239
+ iter_num == 0,
240
+ lambda x: jax.debug.callback(self._define_tqdm, x, ordered=True),
241
+ lambda x: None,
242
+ data
243
+ )
244
+ _ = jax.lax.cond(
245
+ iter_num % self.print_freq == (self.print_freq - 1),
246
+ lambda x: jax.debug.callback(self._update_tqdm, x, ordered=True),
247
+ lambda x: None,
248
+ data
249
+ )
250
+ _ = jax.lax.cond(
251
+ iter_num == self.n - 1,
252
+ lambda x: jax.debug.callback(self._close_tqdm, x, ordered=True),
253
+ lambda x: None,
254
+ data
255
+ )