From 51db29fa192914cf4574981454e98f291e7888d2 Mon Sep 17 00:00:00 2001 From: Stephan Seitz <stephan.seitz@fau.de> Date: Wed, 7 Aug 2019 12:44:54 +0200 Subject: [PATCH] Code cosmetics --- src/pystencils_autodiff/autodiff.py | 91 ++++++++++--------- .../backends/test_torch_native_compilation.py | 1 - 2 files changed, 47 insertions(+), 45 deletions(-) diff --git a/src/pystencils_autodiff/autodiff.py b/src/pystencils_autodiff/autodiff.py index b7e5f5c..8a4e55b 100644 --- a/src/pystencils_autodiff/autodiff.py +++ b/src/pystencils_autodiff/autodiff.py @@ -20,48 +20,6 @@ class DiffModes(str, Enum): TF_MAD = 'transposed-forward' -def _has_exclusive_writes(assignment_collection): - """ - Simple check for exclusive (non-overlapping) writes. - I.e. AssignmentCollection can be executed safely in parallel without caring about race conditions. - No writes on same spatial location (considering all possible shifts). - - The checked condition might be violated if using DiffModes.TRANSPOSED - """ - - assignments = assignment_collection.main_assignments - write_field_accesses = [a.lhs for a in assignments if isinstance(a.lhs, ps.Field.Access)] - - exclusive_writes = set() - for a in write_field_accesses: - - if (a.field, a.index) in exclusive_writes: - return False - else: - exclusive_writes.add((a.field, a.index)) - - return True - - -def get_jacobian_of_assignments(assignments, diff_variables): - """ - Calculates the Jacobian of iterable of assignments wrt. diff_variables - - Arguments: - assignments (List[pystencils.Assignment]): A collection of assignments or a AssignmentCollection - diff_variables (List[sympy.Symbol]): List of variables used to differentiate - - Returns: - sp.Matrix -- Jacobian of statements - """ - - if hasattr(assignments, 'main_assignments'): - assignments = assignments.main_assignments - - rhs = sp.Matrix([e.rhs for e in assignments]) - return rhs.jacobian(diff_variables) - - class AutoDiffOp: _REPR_TEMPLATE = jinja2.Template( """Forward: @@ -588,8 +546,11 @@ def create_backward_assignments(forward_assignments, class AutoDiffAstPair: - """A pair of ASTs of forward and backward kernel. - Just needed, if compilation from AssignmentCollection is not sufficient and you want to manipulate the ASTs""" + """ + A pair of ASTs of forward and backward kernel. + + Just needed, if compilation from AssignmentCollection is not sufficient and you want to manipulate the ASTs + """ def __init__(self, forward_ast, backward_ast, compilation_target='cpu'): self.forward_ast = forward_ast @@ -609,3 +570,45 @@ class AutoDiffAstPair: def __call__(self, *args, **kwargs): return self.forward(*args, **kwargs) + + +def _has_exclusive_writes(assignment_collection): + """ + Simple check for exclusive (non-overlapping) writes. + I.e. AssignmentCollection can be executed safely in parallel without caring about race conditions. + No writes on same spatial location (considering all possible shifts). + + The checked condition might be violated if using DiffModes.TRANSPOSED + """ + + assignments = assignment_collection.main_assignments + write_field_accesses = [a.lhs for a in assignments if isinstance(a.lhs, ps.Field.Access)] + + exclusive_writes = set() + for a in write_field_accesses: + + if (a.field, a.index) in exclusive_writes: + return False + else: + exclusive_writes.add((a.field, a.index)) + + return True + + +def get_jacobian_of_assignments(assignments, diff_variables): + """ + Calculates the Jacobian of iterable of assignments wrt. diff_variables + + Arguments: + assignments (List[pystencils.Assignment]): A collection of assignments or a AssignmentCollection + diff_variables (List[sympy.Symbol]): List of variables used to differentiate + + Returns: + sp.Matrix -- Jacobian of statements + """ + + if hasattr(assignments, 'main_assignments'): + assignments = assignments.main_assignments + + rhs = sp.Matrix([e.rhs for e in assignments]) + return rhs.jacobian(diff_variables) diff --git a/tests/backends/test_torch_native_compilation.py b/tests/backends/test_torch_native_compilation.py index 2371c5e..0b3166c 100644 --- a/tests/backends/test_torch_native_compilation.py +++ b/tests/backends/test_torch_native_compilation.py @@ -44,7 +44,6 @@ def test_jit(): lltm_cuda = CUDAExtension(join(dirname(__file__), 'lltm_cuda'), [cpp_file, cuda_file]) assert lltm_cuda is not None - print('hallo') def test_torch_native_compilation(): -- GitLab