prpy 0.4.2__tar.gz → 0.4.5__tar.gz

{prpy-0.4.2 → prpy-0.4.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: prpy
-Version: 0.4.2
+Version: 0.4.5
 Summary: Collection of Python utils for signal, image, and video processing
 Author-email: Philipp Rouast <philipp@rouast.com>
 License: MIT License
@@ -11,7 +11,7 @@ Classifier: Programming Language :: Python :: 3
 Classifier: Operating System :: OS Independent
 Classifier: Development Status :: 4 - Beta
 Classifier: License :: OSI Approved :: MIT License
-Requires-Python: >=3.10
+Requires-Python: <3.13,>=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: numpy>=1.24
@@ -29,7 +29,7 @@ Requires-Dist: scipy; extra == "numpy-min"
 Requires-Dist: Pillow; extra == "numpy-min"
 Requires-Dist: psutil; extra == "numpy-min"
 Provides-Extra: tensorflow
-Requires-Dist: tensorflow<2.13; extra == "tensorflow"
+Requires-Dist: tensorflow>2.17.0; extra == "tensorflow"
 Provides-Extra: torch
 Requires-Dist: torch>=2.1.0; extra == "torch"
 Provides-Extra: test
@@ -50,8 +50,6 @@ It contains subpackages for working with `numpy`, `ffmpeg`, `tensorflow`, and `t
 
 General prerequisites are `python>=3.10` and `ffmpeg` installed and accessible via the `$PATH` environment variable.
 
-- Please note: If using `numpy` or `tensorflow` options, we only support Python `<3.12` because of the dependencies. 
-
 The easiest way to install the latest version of `prpy`:
 
 ```

{prpy-0.4.2 → prpy-0.4.5}/README.md

@@ -10,8 +10,6 @@ It contains subpackages for working with `numpy`, `ffmpeg`, `tensorflow`, and `t
 
 General prerequisites are `python>=3.10` and `ffmpeg` installed and accessible via the `$PATH` environment variable.
 
-- Please note: If using `numpy` or `tensorflow` options, we only support Python `<3.12` because of the dependencies. 
-
 The easiest way to install the latest version of `prpy`:
 
 ```

{prpy-0.4.2 → prpy-0.4.5}/prpy/ffmpeg/probe.py

@@ -172,9 +172,18 @@ def probe_video(
 
   return fps, total_frames, width, height, codec, bitrate, rotation, issues
 
+def _parse_ffprobe_csv_floats(raw: str) -> np.ndarray:
+  """Parse ffprobe's `-of csv=p=0` stdout (one value per line, possibly with a
+  trailing comma) into a float array."""
+  cleaned = "\n".join(line.rstrip(",") for line in raw.splitlines())
+  return np.fromstring(cleaned, dtype=float, sep="\n")
+
 def probe_video_frame_timestamps(path: str, sanity_check: bool = False) -> list:
   """Probe a video file for a best effort estimate of its frame timestamps.
 
+  Tries reading each packet's pts directly and sorting (no decode needed).
+  Falls back to decoding every frame otherwise.
+
   Args:
     path: The path of the video.
     sanity_check: Whether to check result against ffprobe count
@@ -184,23 +193,45 @@ def probe_video_frame_timestamps(path: str, sanity_check: bool = False) -> list:
   if not os.path.exists(path):
     raise FileNotFoundError(f"File {path} does not exist")
 
-  # Build the ffprobe command to extract the best effort timestamp for each frame.
-  ts_cmd = [
-    "ffprobe", "-v", "error",
-    "-select_streams", "v:0",
-    "-show_entries", "frame=best_effort_timestamp_time",
-    "-of", "csv=p=0",
-    path
-  ]
+  timestamps = None
   try:
-    proc = subprocess.run(ts_cmd, capture_output=True, text=True, check=True)
-    raw = proc.stdout
+    proc = subprocess.run([
+      "ffprobe", "-v", "error",
+      "-select_streams", "v:0",
+      "-show_entries", "packet=pts_time",
+      "-of", "csv=p=0",
+      path
+    ], capture_output=True, text=True, check=True)
+    lines = proc.stdout.splitlines()
+    if lines and all(line.rstrip(",") for line in lines):
+      candidate = np.sort(_parse_ffprobe_csv_floats(proc.stdout))
+      # A handful of tied pts (e.g. a GOP-boundary rounding artifact) is fine
+      # -- sorting already guarantees non-decreasing order, so the only real
+      # question is whether too many packets are tied to trust the result.
+      n_dupes = int(np.sum(np.diff(candidate) <= 0)) if candidate.size else 0
+      if candidate.size == len(lines) and n_dupes <= max(1, int(0.01 * candidate.size)):
+        timestamps = candidate
+      else:
+        logging.debug("Packet-level pts unusable (incomplete or too many duplicates); falling back.")
   except subprocess.CalledProcessError as e:
-    logging.error("ffprobe error: %s", e)
-    return []
-  # Clean trailing commas and process with numpy
-  cleaned = "\n".join(line.rstrip(",") for line in raw.splitlines())
-  timestamps = np.fromstring(cleaned, dtype=float, sep="\n")
+    logging.debug("Fast packet-level probe failed, falling back: %s", e)
+
+  if timestamps is None:
+    # Fallback: decode every frame to read its best-effort timestamp.
+    ts_cmd = [
+      "ffprobe", "-v", "error",
+      "-select_streams", "v:0",
+      "-show_entries", "frame=best_effort_timestamp_time",
+      "-of", "csv=p=0",
+      path
+    ]
+    try:
+      proc = subprocess.run(ts_cmd, capture_output=True, text=True, check=True)
+      raw = proc.stdout
+    except subprocess.CalledProcessError as e:
+      logging.error("ffprobe error: %s", e)
+      return []
+    timestamps = _parse_ffprobe_csv_floats(raw)
 
   if sanity_check:
     # Grab the frame count

{prpy-0.4.2 → prpy-0.4.5}/prpy/numpy/filters.py

@@ -21,6 +21,7 @@
 import math
 import numpy as np
 from scipy import signal
+from scipy.linalg import solveh_banded
 from scipy.ndimage import uniform_filter1d
 from scipy.sparse import spdiags
 from typing import Union, Tuple
@@ -151,18 +152,20 @@ def detrend(
   T = z.shape[axis]
   if T < 3:
    return z
-  # Identity matrix
-  I = np.identity(T)
   # Regularization matrix
   D2 = spdiags(
     [np.ones(T), -2*np.ones(T), np.ones(T)],
-    [0, 1, 2], (T-2), T).toarray()
-  # Inverse of I+lambda^2*D2’*D2
-  A = I + (Lambda**2) * np.dot(D2.T, D2)
+    [0, 1, 2], (T-2), T)
+  # I+lambda^2*D2’*D2 is symmetric pentadiagonal, so solve it in banded form
+  DTD = (D2.T @ D2).todia()
+  ab = np.zeros((3, T))
+  ab[2] = 1. + (Lambda**2) * DTD.diagonal(0)
+  ab[1, 1:] = (Lambda**2) * DTD.diagonal(1)
+  ab[0, 2:] = (Lambda**2) * DTD.diagonal(2)
   # Compute the detrending operation (vectorized)
   if axis == 0:
     z = np.transpose(z)
-  z_trend = np.linalg.solve(A, z.T)
+  z_trend = solveh_banded(ab, z.T)
   proc_z = (z.T - z_trend)
   if axis == 1:
     proc_z = np.transpose(proc_z)

{prpy-0.4.2 → prpy-0.4.5}/prpy/tensorflow/model_saver.py

@@ -56,7 +56,7 @@ class ModelSaver(object):
       dir: str = "checkpoints",
       keep_best: int = 5,
       keep_latest: int = 1,
-      save_format: str = "tf",
+      save_format: str = "keras",
       save_optimizer: bool = False,
       compare_fn: Callable[[float, float], bool] = lambda x,y: x.score < y.score,
       sort_reverse: bool = False,
@@ -68,7 +68,7 @@ class ModelSaver(object):
       dir: The directory where models should be saved
       keep_best: The number of best scoring models to keep
       keep_latest: The number of latest models to keep
-      save_format: Model format for saving ['tf' or 'h5']
+      save_format: Model format for saving ['keras' or 'h5']
       save_optimizer: Also save optimizer state?
       compare_fn: Function that compares two scores
       sort_reverse: Reverse sort order?
@@ -110,15 +110,27 @@ class ModelSaver(object):
     """
     assert isinstance(model, tf.keras.Model)
     assert isinstance(filepath, str)
-    if self.save_format == 'h5': filepath += '.h5'
-    # Save model
     if self.save_optimizer:
-      model.save(
-        filepath=filepath, overwrite=True, include_optimizer=True,
-        save_format=self.save_format)
+      # Full model saving requires .keras or .h5
+      if self.save_format == 'h5': 
+        filepath += '.h5'
+      elif self.save_format == 'keras':
+        filepath += '.keras'
+      else:
+        raise ValueError(f"Unsupported save_format: {self.save_format}. Use 'keras' or 'h5'.")
+      model.save(filepath=filepath, overwrite=True)
     else:
-      model.save_weights(
-        filepath=filepath, overwrite=True, save_format=self.save_format)
+      # Keras 3 strictly demands .weights.h5 for weights-only saves, regardless of format choice
+      filepath += '.weights.h5'
+      model.save_weights(filepath=filepath, overwrite=True)
+
+  def _cleanup_old_files(self, filepath_prefix: str):
+    """Helper to clean up old files or directories robustly."""
+    for file_or_dir in glob.glob(fr"{filepath_prefix}*"):
+      if os.path.isdir(file_or_dir):
+        shutil.rmtree(file_or_dir)
+      else:
+        os.remove(file_or_dir)
 
   def save_keep(
       self,
@@ -170,11 +182,7 @@ class ModelSaver(object):
       self.__save(model, filepath=candidate.filepath)
       # Prune candidate
       for candidate in self.latest_candidates[self.keep_latest:]:
-        for file in glob.glob(fr"{candidate.filepath}*"):
-          if self.save_format == 'tf' and self.save_optimizer:
-            shutil.rmtree(file)
-          else:
-            os.remove(file)
+        self._cleanup_old_files(candidate.filepath)
       self.latest_candidates = self.latest_candidates[0:self.keep_latest]
 
   def save_best(
@@ -211,11 +219,7 @@ class ModelSaver(object):
       # Prune candidates
       for candidate in self.best_candidates[self.keep_best:]:
         self.log_fn(f"Removing old model {candidate.filepath} with score {candidate.score:.4f}")
-        for file in glob.glob(fr"{candidate.filepath}*"):
-          if self.save_format == 'tf' and self.save_optimizer:
-            shutil.rmtree(file)
-          else:
-            os.remove(file)
+        self._cleanup_old_files(candidate.filepath)
       self.best_candidates = self.best_candidates[0:self.keep_best]
     else:
       # Skip the candidate

{prpy-0.4.2 → prpy-0.4.5}/prpy/tensorflow/nan.py

@@ -174,6 +174,15 @@ class ReduceNanSum:
     out = tf.where(tf.reduce_all(tf.logical_not(mask), axis=self.axis),
                    tf.cast(self.default, x.dtype),
                    sum)
+    # TODO(plethnet KNOWN_ISSUES.md #2): this gradient never references
+    # self.weight, even though the forward pass above weights `x` before
+    # summing (`x * weight`). d(out)/dx should be `weight` at finite
+    # positions, not 1. In production this means ReduceNanSum(weight=
+    # loss_weights, ...) (plethnet_tensorflow/engine.py's per-signal loss
+    # combination) has loss_weights affecting only the logged loss value,
+    # not the actual training gradient. See KNOWN_ISSUES.md before fixing -
+    # every plethnet_v* checkpoint trained so far, including production
+    # plethnet_v4, was trained with this behavior.
     def grad(upstream: tf.Tensor) -> tf.Tensor:
       # Tile upstream to match x
       if self.axis is not None:

prpy-0.4.5/prpy/tensorflow/optimizer.py

@@ -0,0 +1,67 @@
+# Copyright (c) 2026 Philipp Rouast
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+import tensorflow as tf
+
+class EpochAdam(tf.keras.optimizers.Adam):
+  """Adam optimizer that natively tracks epochs."""
+  def __init__(self, learning_rate=0.001, **kwargs):
+    super().__init__(learning_rate=learning_rate, **kwargs)
+    with tf.init_scope():
+      self._epochs = tf.Variable(
+        initial_value=0, name="epochs", dtype=tf.int64, trainable=False
+      )
+  @property
+  def epochs(self):
+    return self._epochs
+  @epochs.setter
+  def epochs(self, variable):
+    if self.built:
+      raise RuntimeError(
+        "Cannot set `epochs` to a new Variable after the Optimizer weights "
+        "have been created. Please set `epochs` before calling `apply_gradients()`."
+      )
+    self._epochs = variable 
+  def finish_epoch(self):
+    """Increment epoch count."""
+    self._epochs.assign_add(1)
+
+class EpochAdamW(tf.keras.optimizers.AdamW):
+  """AdamW optimizer that natively tracks epochs."""
+  def __init__(self, learning_rate=0.001, weight_decay=0.004, **kwargs):
+    super().__init__(learning_rate=learning_rate, weight_decay=weight_decay, **kwargs)
+    with tf.init_scope():
+      self._epochs = tf.Variable(
+        initial_value=0, name="epochs", dtype=tf.int64, trainable=False
+      )
+  @property
+  def epochs(self):
+    return self._epochs
+  @epochs.setter
+  def epochs(self, variable):
+    if self.built:
+      raise RuntimeError(
+        "Cannot set `epochs` to a new Variable after the Optimizer weights "
+        "have been created. Please set `epochs` before calling `apply_gradients()`."
+      )
+    self._epochs = variable 
+  def finish_epoch(self):
+    """Increment epoch count."""
+    self._epochs.assign_add(1)

prpy-0.4.5/prpy/torch/loss.py

@@ -0,0 +1,46 @@
+# Copyright (c) 2026 Philipp Rouast
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+import torch
+
+
+def balanced_sample_weights(labels: torch.Tensor, unique: torch.Tensor) -> torch.Tensor:
+  """Calculate weights for a batch of dense categorical labels intended to be
+  multiplied with the losses.
+  - Larger weights for examples of under-represented classes, and smaller
+    weights for overrepresented classes, while keeping the total loss constant.
+  - Important: Only works for dense label representation that equal range from 0 to n.
+  Args:
+    labels: The dense categorical labels of shape (batch_size,) or (batch_size, 1)
+    unique: The unique labels of shape (n_unique_labels,)
+  Returns:
+    weights: The weights with same shape as labels.
+  """
+  original_shape = labels.shape
+  f_labels = labels.reshape(-1).long()
+  minlength = int(unique.max().item()) + 1
+  count = torch.bincount(f_labels, minlength=minlength)[unique.long()]
+  batch_size = f_labels.numel()
+  unique_count = (count > 0).sum()
+  class_weights = torch.where(
+    count > 0, batch_size / count.clamp(min=1).to(torch.float32), torch.zeros_like(count, dtype=torch.float32))
+  class_weights = class_weights / unique_count.to(torch.float32)
+  sample_weights = class_weights[f_labels]
+  return sample_weights.reshape(original_shape)

prpy-0.4.5/prpy/torch/lr_schedule.py

@@ -0,0 +1,54 @@
+# Copyright (c) 2026 Philipp Rouast
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+from typing import List, Union
+
+
+def piecewise_constant_decay_with_warmup(
+    step: Union[int, float],
+    boundaries: List[Union[int, float]],
+    values: List[float],
+    warmup_init_lr: float,
+    warmup_steps: Union[int, float]
+  ) -> float:
+  """Piecewise constant decay with linear warmup.
+  Args:
+    step: The current step (or epoch - unit-agnostic, just compared against
+      boundaries/warmup_steps in whatever unit they're expressed in).
+    boundaries: Step boundaries. len(boundaries) == len(values) - 1.
+    values: The constant LR value for each segment between boundaries.
+    warmup_init_lr: The initial LR at step 0, ramping linearly to values[0]
+      by `warmup_steps`.
+    warmup_steps: Number of steps for the warmup ramp.
+  Returns:
+    The learning rate at `step`.
+  """
+  if len(boundaries) != len(values) - 1:
+    raise ValueError("The length of boundaries should be 1 less than the length of values")
+  if step <= warmup_steps:
+    return warmup_init_lr + step * (values[0] - warmup_init_lr) / warmup_steps
+  if step <= boundaries[0]:
+    return values[0]
+  if step > boundaries[-1]:
+    return values[-1]
+  for low, high, v in zip(boundaries[:-1], boundaries[1:], values[1:-1]):
+    if low < step <= high:
+      return v
+  raise ValueError(f"step={step} did not match any segment of boundaries={boundaries}")

prpy-0.4.5/prpy/torch/nan.py

@@ -0,0 +1,200 @@
+# Copyright (c) 2026 Philipp Rouast
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+from typing import Optional, Tuple, Union
+
+import torch
+
+Dims = Union[int, Tuple[int, ...], None]
+
+
+def _as_dims(dim: Dims) -> Optional[Tuple[int, ...]]:
+  if dim is None:
+    return None
+  return (dim,) if isinstance(dim, int) else tuple(dim)
+
+
+def reduce_nanmean(x: torch.Tensor, dim: Dims = None) -> torch.Tensor:
+  """torch.mean, ignoring non-finite vals.
+  - Returns `nan` for all-nan slices.
+  Args:
+    x: The input tensor.
+    dim: The dimension(s) to reduce.
+  Returns:
+    The reduced tensor.
+  """
+  mask = torch.isfinite(x)
+  safe_x = torch.where(mask, x, torch.zeros_like(x))
+  dims = _as_dims(dim)
+  numerator = safe_x.sum(dim=dims) if dims is not None else safe_x.sum()
+  denominator = (mask.sum(dim=dims) if dims is not None else mask.sum()).to(x.dtype)
+  return numerator / denominator
+
+
+def reduce_nansum(
+    x: torch.Tensor,
+    weight: Optional[torch.Tensor] = None,
+    dim: Dims = None,
+    default: float = float('nan')
+  ) -> torch.Tensor:
+  """torch.sum, weighted by weight, ignoring non-finite values.
+  - Returns default for all-nan slices.
+  Args:
+    x: The input tensor.
+    weight: The weight tensor, with the same shape as x or broadcastable to it.
+    dim: The dimension(s) to reduce.
+    default: The value to return for all-non-finite slices.
+  Returns:
+    The reduced tensor.
+  """
+  mask = torch.isfinite(x)
+  safe_x = torch.where(mask, x, torch.zeros_like(x))
+  if weight is not None:
+    safe_weight = torch.where(mask, weight, torch.zeros_like(weight))
+    safe_x = safe_x * safe_weight.to(x.dtype)
+  dims = _as_dims(dim)
+  total = safe_x.sum(dim=dims) if dims is not None else safe_x.sum()
+  all_nan = (~mask).all(dim=dims) if dims is not None else (~mask).all()
+  return torch.where(all_nan, torch.tensor(default, dtype=x.dtype, device=x.device), total)
+
+
+def _unsqueeze_at(t: torch.Tensor, dims: Tuple[int, ...]) -> torch.Tensor:
+  # Ascending order so earlier insertions don't shift later indices.
+  for d in sorted(dims):
+    t = t.unsqueeze(d)
+  return t
+
+
+class _ReduceNanMeanFn(torch.autograd.Function):
+  @staticmethod
+  def forward(ctx, x, dims):
+    mask = torch.isfinite(x)
+    safe_x = torch.where(mask, x, torch.zeros_like(x))
+    numerator = safe_x.sum(dim=dims) if dims is not None else safe_x.sum()
+    denominator = (mask.sum(dim=dims) if dims is not None else mask.sum()).to(x.dtype)
+    ctx.save_for_backward(mask, denominator)
+    ctx.dims = dims
+    return numerator / denominator
+  @staticmethod
+  def backward(ctx, grad_output):
+    mask, denominator = ctx.saved_tensors
+    if ctx.dims is not None:
+      grad_output = _unsqueeze_at(grad_output, ctx.dims)
+      denominator = _unsqueeze_at(denominator, ctx.dims)
+    dx = torch.where(mask, grad_output / denominator, torch.zeros_like(mask, dtype=grad_output.dtype))
+    return dx, None
+
+
+class ReduceNanMean:
+  """torch.mean, ignoring non-finite values. Supports gradient.
+
+  Behavior when x is non-finite:
+  - out: Non-finite vals in a slice contribute 0
+  - out: All-non-finite slices are nan
+  - grad = 0 at non-finite input positions (including for all-non-finite slices)
+  """
+  def __init__(self, dim: Dims = None):
+    self.dims = _as_dims(dim)
+  def __call__(self, x: torch.Tensor) -> torch.Tensor:
+    return _ReduceNanMeanFn.apply(x, self.dims)
+
+
+class _ReduceNanSumFn(torch.autograd.Function):
+  @staticmethod
+  def forward(ctx, x, weight, dims, default):
+    mask = torch.isfinite(x)
+    safe_x = torch.where(mask, x, torch.zeros_like(x))
+    if weight is not None:
+      safe_weight = torch.where(mask, weight, torch.zeros_like(weight)).to(x.dtype)
+      weighted = safe_x * safe_weight
+    else:
+      weighted = safe_x
+    total = weighted.sum(dim=dims) if dims is not None else weighted.sum()
+    all_nan = (~mask).all(dim=dims) if dims is not None else (~mask).all()
+    out = torch.where(all_nan, torch.tensor(default, dtype=x.dtype, device=x.device), total)
+    ctx.save_for_backward(mask)
+    ctx.dims = dims
+    return out
+  @staticmethod
+  def backward(ctx, grad_output):
+    # Deliberately does not multiply by weight - see ReduceNanSum's docstring.
+    (mask,) = ctx.saved_tensors
+    if ctx.dims is not None:
+      grad_output = _unsqueeze_at(grad_output, ctx.dims)
+    grad_output = torch.broadcast_to(grad_output, mask.shape)
+    dx = torch.where(mask, grad_output, torch.zeros_like(mask, dtype=grad_output.dtype))
+    return dx, None, None, None
+
+
+class ReduceNanSum:
+  """torch.sum, weighted by weight, ignoring non-finite values. Supports gradient.
+
+  Behavior when x is non-finite:
+  - out: Non-finite vals in a slice contribute 0
+  - out: All-non-finite slices are set to default value
+  - grad = 0 at non-finite input positions
+
+  TODO(plethnet KNOWN_ISSUES.md #2): when `weight` is given, it scales the
+  forward sum (`x * weight`) but is NOT applied in the gradient - matches a
+  discrepancy in TF's own hand-written gradient, preserved here verbatim for
+  parity with what plethnet_v4 was actually trained with. Don't fix without
+  sign-off - it would change real training dynamics.
+  """
+  def __init__(self, weight: Optional[torch.Tensor] = None, dim: Dims = None, default: float = float('nan')):
+    self.weight = weight
+    self.dims = _as_dims(dim)
+    self.default = default
+  def __call__(self, x: torch.Tensor) -> torch.Tensor:
+    return _ReduceNanSumFn.apply(x, self.weight, self.dims, self.default)
+
+
+class _NanLinearCombinationFn(torch.autograd.Function):
+  @staticmethod
+  def forward(ctx, x, val_1, val_2):
+    val_1 = torch.broadcast_to(val_1, x.shape)
+    val_2 = torch.broadcast_to(val_2, x.shape)
+    ctx.save_for_backward(x, val_1, val_2)
+    return (1 - x) * val_1 + x * val_2
+  @staticmethod
+  def backward(ctx, grad_output):
+    x, val_1, val_2 = ctx.saved_tensors
+    mask = torch.isfinite(val_1) & torch.isfinite(val_2)
+    grad_x = torch.where(mask, grad_output * (val_2 - val_1), torch.zeros_like(x))
+    grad_val_1 = grad_output * (1 - x)
+    grad_val_2 = grad_output * x
+    return grad_x, grad_val_1, grad_val_2
+
+
+class NanLinearCombination:
+  """Linear combination with one fixed element. Supports gradient.
+  - Calculates (1 - x) * val_1 + x * val_2
+  - Behavior when val_1 or val_2 are non-finite: out = nan, grad wrt x = 0
+    (grad wrt val_1/val_2 are ordinary (1-x)/x)
+  """
+  def __call__(self, x: torch.Tensor, val_1: torch.Tensor, val_2: torch.Tensor) -> torch.Tensor:
+    """Compute the linear combination.
+    Args:
+      x: The combination weight. All elements must be in [0, 1]
+      val_1: The first value used in the linear combination.
+      val_2: The second value used in the linear combination.
+    Returns:
+      out: The computed linear combination, broadcast to x's shape.
+    """
+    return _NanLinearCombinationFn.apply(x, val_1, val_2)

{prpy-0.4.2 → prpy-0.4.5}/prpy.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: prpy
-Version: 0.4.2
+Version: 0.4.5
 Summary: Collection of Python utils for signal, image, and video processing
 Author-email: Philipp Rouast <philipp@rouast.com>
 License: MIT License
@@ -11,7 +11,7 @@ Classifier: Programming Language :: Python :: 3
 Classifier: Operating System :: OS Independent
 Classifier: Development Status :: 4 - Beta
 Classifier: License :: OSI Approved :: MIT License
-Requires-Python: >=3.10
+Requires-Python: <3.13,>=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: numpy>=1.24
@@ -29,7 +29,7 @@ Requires-Dist: scipy; extra == "numpy-min"
 Requires-Dist: Pillow; extra == "numpy-min"
 Requires-Dist: psutil; extra == "numpy-min"
 Provides-Extra: tensorflow
-Requires-Dist: tensorflow<2.13; extra == "tensorflow"
+Requires-Dist: tensorflow>2.17.0; extra == "tensorflow"
 Provides-Extra: torch
 Requires-Dist: torch>=2.1.0; extra == "torch"
 Provides-Extra: test
@@ -50,8 +50,6 @@ It contains subpackages for working with `numpy`, `ffmpeg`, `tensorflow`, and `t
 
 General prerequisites are `python>=3.10` and `ffmpeg` installed and accessible via the `$PATH` environment variable.
 
-- Please note: If using `numpy` or `tensorflow` options, we only support Python `<3.12` because of the dependencies. 
-
 The easiest way to install the latest version of `prpy`:
 
 ```