Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion tests/profiler/plot_memory_timeline.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,8 @@
Reads memory traces from json files and plots them on a single graph.
"""

from __future__ import annotations

import argparse
import json
from dataclasses import dataclass
Expand All @@ -24,7 +26,7 @@ class MemoryFrame:
device_id: int # -1 for CPU, 0+ for CUDA devices

@staticmethod
def from_event(event: dict) -> "MemoryFrame":
def from_event(event: dict) -> MemoryFrame:
args = event["args"]
return MemoryFrame(
timestamp=event["ts"],
Expand Down
6 changes: 4 additions & 2 deletions tests/trajectories/_objectives.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
from __future__ import annotations

from abc import ABC, abstractmethod

import torch
Expand Down Expand Up @@ -77,7 +79,7 @@ def quad(x: Tensor, A: Tensor, u: Tensor) -> Tensor:
return x_minus_u @ A @ x_minus_u

@property
def sps_mapping(self) -> "QuadraticSPSMapping":
def sps_mapping(self) -> QuadraticSPSMapping:
return QuadraticSPSMapping(self.As, self.us)


Expand Down Expand Up @@ -118,7 +120,7 @@ def jacobian(self, x: Tensor) -> Tensor:
return torch.diag(torch.stack([2 * x[0], 2 * x[1]]))

@property
def sps_mapping(self) -> "ElementWiseQuadraticSPSMapping":
def sps_mapping(self) -> ElementWiseQuadraticSPSMapping:
return ElementWiseQuadraticSPSMapping(self.n_values)


Expand Down
6 changes: 4 additions & 2 deletions tests/trajectories/_plotters.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
from __future__ import annotations

from abc import ABC, abstractmethod
from typing import TypeAlias

Expand All @@ -16,7 +18,7 @@ class Plotter(ABC):
def __call__(self, ax: plt.Axes) -> None:
pass

def __add__(self, other: "Plotter") -> "Plotter":
def __add__(self, other: Plotter) -> Plotter:
return MultiPlotter((self, other))


Expand All @@ -30,7 +32,7 @@ def __call__(self, ax: plt.Axes) -> None:
class MultiPlotter(Plotter):
"""Plotter applying several plotters."""

def __init__(self, plotters: tuple["Plotter", ...]) -> None:
def __init__(self, plotters: tuple[Plotter, ...]) -> None:
self.plotters = plotters

def __call__(self, ax: plt.Axes) -> None:
Expand Down
Loading