Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Quality] Split utils.h and utils.cpp #2348

Merged
merged 3 commits into from
Aug 2, 2024
Merged

[Quality] Split utils.h and utils.cpp #2348

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
05ecfd4
init
vmoens Aug 1, 2024
c21ff4d
amend
vmoens Aug 1, 2024
01cec2a
amend
vmoens Aug 2, 2024
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
Next Next commit
init
  • Loading branch information
@vmoens
vmoens committed Aug 1, 2024
commit 05ecfd45ca3ccc0d15471af0fca05c4047105d58
41 changes: 41 additions & 0 deletions torchrl/csrc/utils.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,41 @@
// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is licensed under the MIT license found in the
// LICENSE file in the root directory of this source tree.
// utils.h
#include "utils.h"
#include <iostream>
torch::Tensor safetanh(torch::Tensor input, float eps) {
return SafeTanh::apply(input, eps);
}
torch::Tensor safeatanh(torch::Tensor input, float eps) {
return SafeInvTanh::apply(input, eps);
}
torch::Tensor SafeTanh::forward(torch::autograd::AutogradContext* ctx, torch::Tensor input, float eps) {
auto out = torch::tanh(input);
auto lim = 1.0 - eps;
out = out.clamp(-lim, lim);
ctx->save_for_backward({out});
return out;
}
torch::autograd::tensor_list SafeTanh::backward(torch::autograd::AutogradContext* ctx, torch::autograd::tensor_list grad_outputs) {
auto saved = ctx->get_saved_variables();
auto out = saved[0];
auto go = grad_outputs[0];
auto grad = go * (1 - out * out);
return {grad, torch::Tensor()};
}
torch::Tensor SafeInvTanh::forward(torch::autograd::AutogradContext* ctx, torch::Tensor input, float eps) {
auto lim = 1.0 - eps;
auto intermediate = input.clamp(-lim, lim);
ctx->save_for_backward({intermediate});
auto out = torch::atanh(intermediate);
return out;
}
torch::autograd::tensor_list SafeInvTanh::backward(torch::autograd::AutogradContext* ctx, torch::autograd::tensor_list grad_outputs) {
auto saved = ctx->get_saved_variables();
auto input = saved[0];
auto go = grad_outputs[0];
auto grad = go / (1 - input * input);
return {grad, torch::Tensor()};
}
56 changes: 11 additions & 45 deletions torchrl/csrc/utils.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,58 +2,24 @@
//
// This source code is licensed under the MIT license found in the
// LICENSE file in the root directory of this source tree.
// utils.h

#pragma once

#include <torch/extension.h>
#include <torch/torch.h>

#include <iostream>

using namespace torch::autograd;
torch::Tensor safetanh(torch::Tensor input, float eps = 1e-6);
torch::Tensor safeatanh(torch::Tensor input, float eps = 1e-6);

class SafeTanh : public Function<SafeTanh> {
class SafeTanh : public torch::autograd::Function<SafeTanh> {
public:
static torch::Tensor forward(AutogradContext* ctx, torch::Tensor input,
float eps = 1e-6) {
auto out = torch::tanh(input);
auto lim = 1.0 - eps;
out = out.clamp(-lim, lim);
ctx->save_for_backward({out});
return out;
}

static tensor_list backward(AutogradContext* ctx, tensor_list grad_outputs) {
auto saved = ctx->get_saved_variables();
auto out = saved[0];
auto go = grad_outputs[0];
auto grad = go * (1 - out * out);
return {grad, torch::Tensor()};
}
static torch::Tensor forward(torch::autograd::AutogradContext* ctx, torch::Tensor input, float eps);
static torch::autograd::tensor_list backward(torch::autograd::AutogradContext* ctx, torch::autograd::tensor_list grad_outputs);
};

torch::Tensor safetanh(torch::Tensor input, float eps = 1e-6) {
return SafeTanh::apply(input, eps);
}

class SafeInvTanh : public Function<SafeInvTanh> {
class SafeInvTanh : public torch::autograd::Function<SafeInvTanh> {
public:
static torch::Tensor forward(AutogradContext* ctx, torch::Tensor input,
float eps = 1e-6) {
auto lim = 1.0 - eps;
auto intermediate = input.clamp(-lim, lim);
ctx->save_for_backward({intermediate});
auto out = torch::atanh(intermediate);
return out;
}

static tensor_list backward(AutogradContext* ctx, tensor_list grad_outputs) {
auto saved = ctx->get_saved_variables();
auto input = saved[0];
auto go = grad_outputs[0];
auto grad = go / (1 - input * input);
return {grad, torch::Tensor()};
}
static torch::Tensor forward(torch::autograd::AutogradContext* ctx, torch::Tensor input, float eps);
static torch::autograd::tensor_list backward(torch::autograd::AutogradContext* ctx, torch::autograd::tensor_list grad_outputs);
};

torch::Tensor safeatanh(torch::Tensor input, float eps = 1e-6) {
return SafeInvTanh::apply(input, eps);
}
Loading