Coverage for muutils / tensor_utils.py: 99%

1"""utilities for working with tensors and arrays. 

2 

3notably: 

4 

5- `TYPE_TO_JAX_DTYPE` : a mapping from python, numpy, and torch types to `jaxtyping` types 

6- `DTYPE_MAP` mapping string representations of types to their type 

7- `TORCH_DTYPE_MAP` mapping string representations of types to torch types 

8- `compare_state_dicts` for comparing two state dicts and giving a detailed error message on whether if was keys, shapes, or values that didn't match 

9 

10""" 

11 

12from __future__ import annotations 

13 

14import json 

15import typing 

16from typing import Any 

17 

18import jaxtyping 

19import numpy as np 

20import torch 

21 

22from muutils.dictmagic import dotlist_to_nested_dict 

23 

24# pylint: disable=missing-class-docstring 

25 

26 

27TYPE_TO_JAX_DTYPE: dict[Any, Any] = { 

28 float: jaxtyping.Float, 

29 int: jaxtyping.Int, 

30 jaxtyping.Float: jaxtyping.Float, 

31 jaxtyping.Int: jaxtyping.Int, 

32 # bool 

33 bool: jaxtyping.Bool, 

34 jaxtyping.Bool: jaxtyping.Bool, 

35 np.bool_: jaxtyping.Bool, 

36 torch.bool: jaxtyping.Bool, 

37 # numpy float 

38 np.float16: jaxtyping.Float, 

39 np.float32: jaxtyping.Float, 

40 np.float64: jaxtyping.Float, 

41 np.half: jaxtyping.Float, 

42 np.single: jaxtyping.Float, 

43 np.double: jaxtyping.Float, 

44 # numpy int 

45 np.int8: jaxtyping.Int, 

46 np.int16: jaxtyping.Int, 

47 np.int32: jaxtyping.Int, 

48 np.int64: jaxtyping.Int, 

49 np.longlong: jaxtyping.Int, 

50 np.short: jaxtyping.Int, 

51 np.uint8: jaxtyping.Int, 

52 # torch float 

53 torch.float: jaxtyping.Float, 

54 torch.float16: jaxtyping.Float, 

55 torch.float32: jaxtyping.Float, 

56 torch.float64: jaxtyping.Float, 

57 torch.half: jaxtyping.Float, 

58 torch.double: jaxtyping.Float, 

59 torch.bfloat16: jaxtyping.Float, 

60 # torch int 

61 torch.int: jaxtyping.Int, 

62 torch.int8: jaxtyping.Int, 

63 torch.int16: jaxtyping.Int, 

64 torch.int32: jaxtyping.Int, 

65 torch.int64: jaxtyping.Int, 

66 torch.long: jaxtyping.Int, 

67 torch.short: jaxtyping.Int, 

68} 

69"dict mapping python, numpy, and torch types to `jaxtyping` types" 

70 

71# np.float_ and np.int_ were deprecated in numpy 1.20 and removed in 2.0 

72# use try/except for backwards compatibility and type checker friendliness 

73try: 

74 TYPE_TO_JAX_DTYPE[np.float_] = jaxtyping.Float # type: ignore[attr-defined] # pyright: ignore[reportAttributeAccessIssue] 

75 TYPE_TO_JAX_DTYPE[np.int_] = jaxtyping.Int # type: ignore[attr-defined] # pyright: ignore[reportAttributeAccessIssue] 

76except AttributeError: 

77 pass # numpy 2.0+ removed these deprecated aliases 

78 

79 

80# TODO: add proper type annotations to this signature 

81# TODO: maybe get rid of this altogether? 

82# def jaxtype_factory( 

83# name: str, 

84# array_type: type, 

85# default_jax_dtype: type[jaxtyping.Float | jaxtyping.Int | jaxtyping.Bool] = jaxtyping.Float, 

86# legacy_mode: typing.Union[ErrorMode, str] = ErrorMode.WARN, 

87# ) -> type: 

88# """usage: 

89# ``` 

90# ATensor = jaxtype_factory("ATensor", torch.Tensor, jaxtyping.Float) 

91# x: ATensor["dim1 dim2", np.float32] 

92# ``` 

93# """ 

94# legacy_mode_ = ErrorMode.from_any(legacy_mode) 

95 

96# class _BaseArray: 

97# """jaxtyping shorthand 

98# (backwards compatible with older versions of muutils.tensor_utils) 

99 

100# default_jax_dtype = {default_jax_dtype} 

101# array_type = {array_type} 

102# """ 

103 

104# def __new__(cls, *args: Any, **kwargs: Any) -> typing.NoReturn: 

105# raise TypeError("Type FArray cannot be instantiated.") 

106 

107# def __init_subclass__(cls, *args: Any, **kwargs: Any) -> typing.NoReturn: 

108# raise TypeError(f"Cannot subclass {cls.__name__}") 

109 

110# @classmethod 

111# def param_info(cls, params: typing.Union[str, tuple[Any, ...]]) -> str: 

112# """useful for error printing""" 

113# return "\n".join( 

114# f"{k} = {v}" 

115# for k, v in { 

116# "cls.__name__": cls.__name__, 

117# "cls.__doc__": cls.__doc__, 

118# "params": params, 

119# "type(params)": type(params), 

120# }.items() 

121# ) 

122 

123# @typing._tp_cache # type: ignore[attr-defined] # pyright: ignore[reportUnknownMemberType, reportAttributeAccessIssue] 

124# def __class_getitem__(cls, params: typing.Union[str, tuple[Any, ...]]) -> type: # type: ignore[misc] 

125# # MyTensor["dim1 dim2"] 

126# if isinstance(params, str): 

127# return default_jax_dtype[array_type, params] 

128 

129# elif isinstance(params, tuple): 

130# if len(params) != 2: 

131# raise Exception( 

132# f"unexpected type for params, expected tuple of length 2 here:\n{cls.param_info(params)}" 

133# ) 

134 

135# if isinstance(params[0], str): 

136# # MyTensor["dim1 dim2", int] 

137# return TYPE_TO_JAX_DTYPE[params[1]][array_type, params[0]] 

138 

139# elif isinstance(params[0], tuple): 

140# legacy_mode_.process( 

141# f"legacy type annotation was used:\n{cls.param_info(params) = }", 

142# except_cls=Exception, 

143# ) 

144# # MyTensor[("dim1", "dim2"), int] 

145# shape_anot: list[str] = list() 

146# for x in params[0]: 

147# if isinstance(x, str): 

148# shape_anot.append(x) 

149# elif isinstance(x, int): 

150# shape_anot.append(str(x)) 

151# elif isinstance(x, tuple): 

152# shape_anot.append("".join(str(y) for y in x)) 

153# else: 

154# raise Exception( 

155# f"unexpected type for params, expected first part to be str, int, or tuple:\n{cls.param_info(params)}" 

156# ) 

157 

158# return TYPE_TO_JAX_DTYPE[params[1]][ 

159# array_type, " ".join(shape_anot) 

160# ] 

161# else: 

162# raise Exception( 

163# f"unexpected type for params:\n{cls.param_info(params)}" 

164# ) 

165 

166# _BaseArray.__name__ = name 

167 

168# if _BaseArray.__doc__ is None: 

169# _BaseArray.__doc__ = "{default_jax_dtype = }\n{array_type = }" 

170 

171# _BaseArray.__doc__ = _BaseArray.__doc__.format( 

172# default_jax_dtype=repr(default_jax_dtype), 

173# array_type=repr(array_type), 

174# ) 

175 

176# return _BaseArray 

177 

178 

179if typing.TYPE_CHECKING: 

180 # these class definitions are only used here to make pylint happy, 

181 # but they make mypy unhappy and there is no way to only run if not mypy 

182 # so, later on we have more ignores 

183 class ATensor(torch.Tensor): 

184 @typing._tp_cache # type: ignore[attr-defined] # pyright: ignore[reportUnknownMemberType, reportAttributeAccessIssue] 

185 def __class_getitem__(cls, params: typing.Union[str, tuple[Any, ...]]) -> type: 

186 raise NotImplementedError() 

187 

188 class NDArray(torch.Tensor): 

189 @typing._tp_cache # type: ignore[attr-defined] # pyright: ignore[reportUnknownMemberType, reportAttributeAccessIssue] 

190 def __class_getitem__(cls, params: typing.Union[str, tuple[Any, ...]]) -> type: 

191 raise NotImplementedError() 

192 

193 

194# ATensor = jaxtype_factory("ATensor", torch.Tensor, jaxtyping.Float) # type: ignore[misc, assignment] 

195 

196# NDArray = jaxtype_factory("NDArray", np.ndarray, jaxtyping.Float) # type: ignore[misc, assignment] 

197 

198 

199def numpy_to_torch_dtype(dtype: typing.Union[np.dtype, torch.dtype]) -> torch.dtype: 

200 """convert numpy dtype to torch dtype""" 

201 if isinstance(dtype, torch.dtype): 

202 return dtype 

203 else: 

204 return torch.from_numpy(np.array(0, dtype=dtype)).dtype 

205 

206 

207DTYPE_LIST: list[Any] = [ 

208 *[ 

209 bool, 

210 int, 

211 float, 

212 ], 

213 *[ 

214 # ---------- 

215 # pytorch 

216 # ---------- 

217 # floats 

218 torch.float, 

219 torch.float32, 

220 torch.float64, 

221 torch.half, 

222 torch.double, 

223 torch.bfloat16, 

224 # complex 

225 torch.complex64, 

226 torch.complex128, 

227 # ints 

228 torch.int, 

229 torch.int8, 

230 torch.int16, 

231 torch.int32, 

232 torch.int64, 

233 torch.long, 

234 torch.short, 

235 # simplest 

236 torch.uint8, 

237 torch.bool, 

238 ], 

239 *[ 

240 # ---------- 

241 # numpy 

242 # ---------- 

243 # floats 

244 np.float16, 

245 np.float32, 

246 np.float64, 

247 np.half, 

248 np.single, 

249 np.double, 

250 # complex 

251 np.complex64, 

252 np.complex128, 

253 # ints 

254 np.int8, 

255 np.int16, 

256 np.int32, 

257 np.int64, 

258 np.longlong, 

259 np.short, 

260 # simplest 

261 np.uint8, 

262 np.bool_, 

263 ], 

264] 

265"list of all the python, numpy, and torch numerical types I could think of" 

266 

267# np.float_ and np.int_ were deprecated in numpy 1.20 and removed in 2.0 

268try: 

269 DTYPE_LIST.extend([np.float_, np.int_]) # type: ignore[attr-defined] # pyright: ignore[reportAttributeAccessIssue] 

270except AttributeError: 

271 pass # numpy 2.0+ removed these deprecated aliases 

272 

273DTYPE_MAP: dict[str, Any] = { 

274 **{str(x): x for x in DTYPE_LIST}, 

275 **{dtype.__name__: dtype for dtype in DTYPE_LIST if dtype.__module__ == "numpy"}, 

276} 

277"mapping from string representations of types to their type" 

278 

279TORCH_DTYPE_MAP: dict[str, torch.dtype] = { 

280 key: numpy_to_torch_dtype(dtype) for key, dtype in DTYPE_MAP.items() 

281} 

282"mapping from string representations of types to specifically torch types" 

283 

284# no idea why we have to do this, smh 

285DTYPE_MAP["bool"] = np.bool_ 

286TORCH_DTYPE_MAP["bool"] = torch.bool 

287 

288 

289TORCH_OPTIMIZERS_MAP: dict[str, typing.Type[torch.optim.Optimizer]] = { 

290 "Adagrad": torch.optim.Adagrad, 

291 "Adam": torch.optim.Adam, 

292 "AdamW": torch.optim.AdamW, 

293 "SparseAdam": torch.optim.SparseAdam, 

294 "Adamax": torch.optim.Adamax, 

295 "ASGD": torch.optim.ASGD, 

296 "LBFGS": torch.optim.LBFGS, 

297 "NAdam": torch.optim.NAdam, 

298 "RAdam": torch.optim.RAdam, 

299 "RMSprop": torch.optim.RMSprop, 

300 "Rprop": torch.optim.Rprop, 

301 "SGD": torch.optim.SGD, 

302} 

303 

304 

305def pad_tensor( 

306 tensor: jaxtyping.Shaped[torch.Tensor, "dim1"], # noqa: F821 

307 padded_length: int, 

308 pad_value: float = 0.0, 

309 rpad: bool = False, 

310) -> jaxtyping.Shaped[torch.Tensor, "padded_length"]: # noqa: F821 

311 """pad a 1-d tensor on the left with pad_value to length `padded_length` 

312 

313 set `rpad = True` to pad on the right instead""" 

314 

315 temp: list[torch.Tensor] = [ 

316 torch.full( 

317 (padded_length - tensor.shape[0],), 

318 pad_value, 

319 dtype=tensor.dtype, 

320 device=tensor.device, 

321 ), 

322 tensor, 

323 ] 

324 

325 if rpad: 

326 temp.reverse() 

327 

328 return torch.cat(temp) 

329 

330 

331def lpad_tensor( 

332 tensor: torch.Tensor, padded_length: int, pad_value: float = 0.0 

333) -> torch.Tensor: 

334 """pad a 1-d tensor on the left with pad_value to length `padded_length`""" 

335 return pad_tensor(tensor, padded_length, pad_value, rpad=False) 

336 

337 

338def rpad_tensor( 

339 tensor: torch.Tensor, pad_length: int, pad_value: float = 0.0 

340) -> torch.Tensor: 

341 """pad a 1-d tensor on the right with pad_value to length `pad_length`""" 

342 return pad_tensor(tensor, pad_length, pad_value, rpad=True) 

343 

344 

345def pad_array( 

346 array: jaxtyping.Shaped[np.ndarray, "dim1"], # noqa: F821 

347 padded_length: int, 

348 pad_value: float = 0.0, 

349 rpad: bool = False, 

350) -> jaxtyping.Shaped[np.ndarray, "padded_length"]: # noqa: F821 

351 """pad a 1-d array on the left with pad_value to length `padded_length` 

352 

353 set `rpad = True` to pad on the right instead""" 

354 

355 temp: list[np.ndarray] = [ 

356 np.full( 

357 (padded_length - array.shape[0],), 

358 pad_value, 

359 dtype=array.dtype, 

360 ), 

361 array, 

362 ] 

363 

364 if rpad: 

365 temp.reverse() 

366 

367 return np.concatenate(temp) 

368 

369 

370def lpad_array( 

371 array: np.ndarray, padded_length: int, pad_value: float = 0.0 

372) -> np.ndarray: 

373 """pad a 1-d array on the left with pad_value to length `padded_length`""" 

374 return pad_array(array, padded_length, pad_value, rpad=False) 

375 

376 

377def rpad_array( 

378 array: np.ndarray, pad_length: int, pad_value: float = 0.0 

379) -> np.ndarray: 

380 """pad a 1-d array on the right with pad_value to length `pad_length`""" 

381 return pad_array(array, pad_length, pad_value, rpad=True) 

382 

383 

384def get_dict_shapes(d: dict[str, "torch.Tensor"]) -> dict[str, tuple[int, ...]]: 

385 """given a state dict or cache dict, compute the shapes and put them in a nested dict""" 

386 return dotlist_to_nested_dict({k: tuple(v.shape) for k, v in d.items()}) 

387 

388 

389def string_dict_shapes(d: dict[str, "torch.Tensor"]) -> str: 

390 """printable version of get_dict_shapes""" 

391 return json.dumps( 

392 dotlist_to_nested_dict( 

393 { 

394 k: str( 

395 tuple(v.shape) 

396 ) # to string, since indent wont play nice with tuples 

397 for k, v in d.items() 

398 } 

399 ), 

400 indent=2, 

401 ) 

402 

403 

404class StateDictCompareError(AssertionError): 

405 """raised when state dicts don't match""" 

406 

407 pass 

408 

409 

410class StateDictKeysError(StateDictCompareError): 

411 """raised when state dict keys don't match""" 

412 

413 pass 

414 

415 

416class StateDictShapeError(StateDictCompareError): 

417 """raised when state dict shapes don't match""" 

418 

419 pass 

420 

421 

422class StateDictValueError(StateDictCompareError): 

423 """raised when state dict values don't match""" 

424 

425 pass 

426 

427 

428def compare_state_dicts( 

429 d1: dict[str, Any], 

430 d2: dict[str, Any], 

431 rtol: float = 1e-5, 

432 atol: float = 1e-8, 

433 verbose: bool = True, 

434) -> None: 

435 """compare two dicts of tensors 

436 

437 # Parameters: 

438 

439 - `d1 : dict` 

440 - `d2 : dict` 

441 - `rtol : float` 

442 (defaults to `1e-5`) 

443 - `atol : float` 

444 (defaults to `1e-8`) 

445 - `verbose : bool` 

446 (defaults to `True`) 

447 

448 # Raises: 

449 

450 - `StateDictKeysError` : keys don't match 

451 - `StateDictShapeError` : shapes don't match (but keys do) 

452 - `StateDictValueError` : values don't match (but keys and shapes do) 

453 """ 

454 # check keys match 

455 d1_keys: set[str] = set(d1.keys()) 

456 d2_keys: set[str] = set(d2.keys()) 

457 symmetric_diff: set[str] = set.symmetric_difference(d1_keys, d2_keys) 

458 keys_diff_1: set[str] = d1_keys - d2_keys 

459 keys_diff_2: set[str] = d2_keys - d1_keys 

460 # sort sets for easier debugging 

461 symmetric_diff = set(sorted(symmetric_diff)) 

462 keys_diff_1 = set(sorted(keys_diff_1)) 

463 keys_diff_2 = set(sorted(keys_diff_2)) 

464 diff_shapes_1: str = ( 

465 string_dict_shapes({k: d1[k] for k in keys_diff_1}) 

466 if verbose 

467 else "(verbose = False)" 

468 ) 

469 diff_shapes_2: str = ( 

470 string_dict_shapes({k: d2[k] for k in keys_diff_2}) 

471 if verbose 

472 else "(verbose = False)" 

473 ) 

474 if not len(symmetric_diff) == 0: 

475 raise StateDictKeysError( 

476 f"state dicts do not match:\n{symmetric_diff = }\n{keys_diff_1 = }\n{keys_diff_2 = }\nd1_shapes = {diff_shapes_1}\nd2_shapes = {diff_shapes_2}" 

477 ) 

478 

479 # check tensors match 

480 shape_failed: list[str] = list() 

481 vals_failed: list[str] = list() 

482 for k, v1 in d1.items(): 

483 v2 = d2[k] 

484 # check shapes first 

485 if not v1.shape == v2.shape: 

486 shape_failed.append(k) 

487 else: 

488 # if shapes match, check values 

489 if not torch.allclose(v1, v2, rtol=rtol, atol=atol): 

490 vals_failed.append(k) 

491 

492 str_shape_failed: str = ( 

493 string_dict_shapes({k: d1[k] for k in shape_failed}) if verbose else "" 

494 ) 

495 str_vals_failed: str = ( 

496 string_dict_shapes({k: d1[k] for k in vals_failed}) if verbose else "" 

497 ) 

498 

499 if not len(shape_failed) == 0: 

500 raise StateDictShapeError( 

501 f"{len(shape_failed)} / {len(d1)} state dict elements don't match in shape:\n{shape_failed = }\n{str_shape_failed}" 

502 ) 

503 if not len(vals_failed) == 0: 

504 raise StateDictValueError( 

505 f"{len(vals_failed)} / {len(d1)} state dict elements don't match in values:\n{vals_failed = }\n{str_vals_failed}" 

506 )